import ast
import unittest

from analyzer.scanner import QGISASTVisitor


class TestVulnerability(unittest.TestCase):
    """Unit tests for Security & Safety audit rules."""

    def test_unsafe_subprocess_shell_true(self):
        code = "subprocess.run('ls -la', shell=True)"
        tree = ast.parse(code)
        visitor = QGISASTVisitor("core.py")
        visitor.visit(tree)

        issues = [i for i in visitor.issues if i["type"] == "UNSAFE_SUBPROCESS"]
        self.assertEqual(len(issues), 1)
        self.assertIn("shell=True", issues[0]["message"])

    def test_unsafe_subprocess_injection(self):
        # f-string injection
        code_f = "subprocess.run(f'ls {user_input}')"
        visitor_f = QGISASTVisitor("core.py")
        visitor_f.visit(ast.parse(code_f))
        issues_f = [i for i in visitor_f.issues if i["type"] == "UNSAFE_SUBPROCESS"]
        self.assertEqual(len(issues_f), 1)
        self.assertIn("injection", issues_f[0]["message"])

        # % operator injection
        code_p = "subprocess.call('ls %s' % user_input)"
        visitor_p = QGISASTVisitor("core.py")
        visitor_p.visit(ast.parse(code_p))
        issues_p = [i for i in visitor_p.issues if i["type"] == "UNSAFE_SUBPROCESS"]
        self.assertEqual(len(issues_p), 1)

    def test_blocking_network_in_ui(self):
        code = "requests.get('https://google.com')"

        # In a UI file
        visitor_ui = QGISASTVisitor("src/gui/dialog.py")
        visitor_ui.visit(ast.parse(code))
        issues_ui = [i for i in visitor_ui.issues if i["type"] == "BLOCKING_NETWORK_CALL"]
        self.assertEqual(len(issues_ui), 1)
        self.assertIn("UI file", issues_ui[0]["message"])

        # In a core file (should not warn about blocking UI)
        visitor_core = QGISASTVisitor("src/core/logic.py")
        visitor_core.visit(ast.parse(code))
        issues_core = [i for i in visitor_core.issues if i["type"] == "BLOCKING_NETWORK_CALL"]
        self.assertEqual(len(issues_core), 0)

    def test_urllib_blocking_in_ui(self):
        code = "urllib.request.urlopen('https://google.com')"
        visitor = QGISASTVisitor("src/ui/main.py")
        visitor.visit(ast.parse(code))
        issues = [i for i in visitor.issues if i["type"] == "BLOCKING_NETWORK_CALL"]
        self.assertEqual(len(issues), 1)


if __name__ == "__main__":
    unittest.main()
