mirror of
https://github.com/Icex0/wp2shell-poc
synced 2026-07-18 19:05:08 +00:00
87 lines
2.4 KiB
Python
87 lines
2.4 KiB
Python
import contextlib
|
|
import io
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from wp2shell import cli
|
|
|
|
|
|
class ShellCommandTests(unittest.TestCase):
|
|
def test_cleanup_runs_when_command_execution_raises(self):
|
|
class FakeSession:
|
|
instance = None
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self.cleaned = False
|
|
type(self).instance = self
|
|
|
|
def login(self, username, password):
|
|
return True
|
|
|
|
def deploy_webshell(self):
|
|
return "/wp-content/plugins/wp2shell_test/wp2shell_test.php"
|
|
|
|
def run(self, path, command):
|
|
raise OSError("connection dropped")
|
|
|
|
def cleanup(self, path):
|
|
self.cleaned = True
|
|
return True
|
|
|
|
with mock.patch.object(cli, "AdminSession", FakeSession):
|
|
rc = cli.main(
|
|
[
|
|
"shell",
|
|
"http://target",
|
|
"--user",
|
|
"admin",
|
|
"--password",
|
|
"password",
|
|
"--cmd",
|
|
"id",
|
|
"--cleanup",
|
|
]
|
|
)
|
|
|
|
self.assertEqual(rc, 1)
|
|
self.assertTrue(FakeSession.instance.cleaned)
|
|
|
|
def test_keep_and_cleanup_are_mutually_exclusive(self):
|
|
parser = cli.build_parser()
|
|
with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
|
|
parser.parse_args(
|
|
[
|
|
"shell",
|
|
"http://target",
|
|
"--user",
|
|
"admin",
|
|
"--password",
|
|
"password",
|
|
"--cmd",
|
|
"id",
|
|
"--keep",
|
|
"--cleanup",
|
|
]
|
|
)
|
|
|
|
def test_rest_route_is_not_a_shell_option(self):
|
|
parser = cli.build_parser()
|
|
with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
|
|
parser.parse_args(
|
|
[
|
|
"shell",
|
|
"http://target",
|
|
"--user",
|
|
"admin",
|
|
"--password",
|
|
"password",
|
|
"--cmd",
|
|
"id",
|
|
"--rest-route",
|
|
]
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|