Files
Icex0-wp2shell-poc/tests/test_cli.py
T
2026-07-18 12:23:05 +12:00

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()