Fix webshell cleanup and CLI behavior

tests/test_cli.py

@@ -0,0 +1,86 @@
+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()

tests/test_shell.py

@@ -0,0 +1,37 @@
+import io
+import unittest
+import urllib.error
+import zipfile
+from unittest import mock
+
+from wp2shell.shell import AdminSession
+
+
+class AdminSessionTests(unittest.TestCase):
+    def test_webshell_changes_to_its_plugin_directory(self):
+        session = AdminSession("http://target")
+        with zipfile.ZipFile(io.BytesIO(session._plugin_zip())) as archive:
+            php = archive.read(f"{session._slug}/{session._slug}.php").decode()
+
+        self.assertIn("chdir(__DIR__);", php)
+
+    def test_cleanup_confirms_a_missing_webshell(self):
+        session = AdminSession("http://target")
+        session.run = mock.Mock(return_value="")
+        session._get = mock.Mock(
+            side_effect=urllib.error.HTTPError("http://target/shell", 404, "missing", {}, None)
+        )
+
+        self.assertTrue(session.cleanup("/shell"))
+        command = session.run.call_args.args[1]
+        self.assertIn("*/wp-content/plugins/*", command)
+
+    def test_cleanup_handles_a_request_failure(self):
+        session = AdminSession("http://target")
+        session.run = mock.Mock(side_effect=OSError("connection dropped"))
+
+        self.assertFalse(session.cleanup("/shell"))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_sqli.py

@@ -0,0 +1,23 @@
+import unittest
+from unittest import mock
+
+from wp2shell.sqli import BlindSQLi
+
+
+class BlindSQLiIntegerTests(unittest.TestCase):
+    def test_integer_rejects_failed_extraction(self):
+        sqli = BlindSQLi(mock.Mock())
+        sqli.extract = mock.Mock(return_value="")
+
+        with self.assertRaisesRegex(ValueError, "expected an integer"):
+            sqli.integer("SELECT COUNT(*)")
+
+    def test_integer_accepts_signed_values(self):
+        sqli = BlindSQLi(mock.Mock())
+        sqli.extract = mock.Mock(return_value=" -12 ")
+
+        self.assertEqual(sqli.integer("SELECT -12"), -12)
+
+
+if __name__ == "__main__":
+    unittest.main()