# security_scanner.py — heuristic code scanner for TDX BLACK HOST
import re
import zipfile
import logging
from pathlib import Path
from typing import Tuple, List

logger = logging.getLogger('TDX.Scanner')

# Heuristic patterns — NOT a perfect malware detector
PYTHON_PATTERNS: List[Tuple[str, str, str]] = [
    (r'\bos\.system\s*\(', 'os.system() call', 'HIGH'),
    (r'\bsubprocess\b', 'subprocess usage', 'HIGH'),
    (r'\beval\s*\(', 'eval() usage', 'HIGH'),
    (r'\bexec\s*\(', 'exec() usage', 'HIGH'),
    (r'\bpickle\.loads?\b', 'pickle deserialization', 'HIGH'),
    (r'\bmarshal\.loads?\b', 'marshal deserialization', 'HIGH'),
    (r'base64\.b64decode', 'base64 decode', 'WARNING'),
    (r'requests\.get\s*\(.*token', 'possible credential exfil', 'HIGH'),
    (r'\.env\b|os\.environ', 'env var access', 'WARNING'),
    (r'__import__\s*\(', 'dynamic import', 'WARNING'),
    (r'compile\s*\(.*exec', 'compile+exec pattern', 'HIGH'),
    (r'socket\.connect\s*\(', 'raw socket connection', 'WARNING'),
    (r'ctypes', 'ctypes usage', 'HIGH'),
    (r'shutil\.rmtree', 'recursive delete', 'WARNING'),
    (r'open\s*\(.*["\']w["\']', 'file write operation', 'WARNING'),
]

JS_PATTERNS: List[Tuple[str, str, str]] = [
    (r'\bchild_process\b', 'child_process module', 'HIGH'),
    (r'\bexec\s*\(', 'exec() call', 'HIGH'),
    (r'\bspawn\s*\(', 'spawn() call', 'HIGH'),
    (r'\beval\s*\(', 'eval() usage', 'HIGH'),
    (r'new Function\s*\(', 'Function constructor', 'HIGH'),
    (r'require\s*\(\s*["\']fs["\']', 'filesystem access', 'WARNING'),
    (r'require\s*\(\s*["\']net["\']', 'raw network access', 'WARNING'),
    (r'process\.env', 'env var access', 'WARNING'),
    (r'Buffer\.from\b', 'Buffer usage', 'WARNING'),
    (r'axios|fetch|http\.request', 'network request', 'WARNING'),
    (r'token|secret|password|api_key', 'credential keyword', 'WARNING'),
]


def _classify(findings: List[Tuple[str, str]]) -> str:
    if not findings:
        return 'SAFE'
    severities = [f[1] for f in findings]
    if 'HIGH' in severities:
        return 'HIGH RISK'
    return 'WARNING'


def scan_python(code: str) -> Tuple[str, List[Tuple[str, str]]]:
    findings = []
    for pattern, label, severity in PYTHON_PATTERNS:
        if re.search(pattern, code, re.IGNORECASE):
            findings.append((label, severity))
    level = _classify(findings)
    return level, findings


def scan_javascript(code: str) -> Tuple[str, List[Tuple[str, str]]]:
    findings = []
    for pattern, label, severity in JS_PATTERNS:
        if re.search(pattern, code, re.IGNORECASE):
            findings.append((label, severity))
    level = _classify(findings)
    return level, findings


def scan_file(file_path: str) -> Tuple[str, str, str]:
    """
    Returns (syntax_ok, risk_level, details_text)
    syntax_ok: 'OK' or error string
    risk_level: 'SAFE' / 'WARNING' / 'HIGH RISK'
    details_text: human-readable findings
    """
    path = Path(file_path)
    suffix = path.suffix.lower()

    try:
        content = path.read_text(encoding='utf-8', errors='replace')
    except Exception as e:
        return f'READ ERROR: {e}', 'HIGH RISK', 'Could not read file.'

    syntax_status = 'OK'

    if suffix == '.py':
        import ast
        try:
            ast.parse(content)
        except SyntaxError as e:
            syntax_status = f'SyntaxError line {e.lineno}: {e.msg}'
        risk, findings = scan_python(content)

    elif suffix == '.js':
        risk, findings = scan_javascript(content)

    elif suffix == '.zip':
        return scan_zip(file_path)

    else:
        return 'N/A', 'WARNING', 'Unknown file type — manual review recommended.'

    details = '\n'.join(
        f'  [{sev}] {label}' for label, sev in findings
    ) if findings else '  No suspicious patterns detected.'

    return syntax_status, risk, details


def scan_zip(zip_path: str) -> Tuple[str, str, str]:
    all_findings = []
    syntax_note = 'OK'
    try:
        with zipfile.ZipFile(zip_path, 'r') as zf:
            for name in zf.namelist():
                if '..' in name or name.startswith('/'):
                    all_findings.append(('Path traversal attempt in ZIP', 'HIGH'))
                    continue
                suffix = Path(name).suffix.lower()
                if suffix in ('.py', '.js'):
                    try:
                        code = zf.read(name).decode('utf-8', errors='replace')
                        if suffix == '.py':
                            _, findings = scan_python(code)
                        else:
                            _, findings = scan_javascript(code)
                        for label, sev in findings:
                            all_findings.append((f'{name}: {label}', sev))
                    except Exception:
                        pass
    except zipfile.BadZipFile:
        return 'Invalid ZIP', 'HIGH RISK', 'File is not a valid ZIP archive.'

    risk = _classify(all_findings)
    details = '\n'.join(
        f'  [{sev}] {label}' for label, sev in all_findings
    ) if all_findings else '  No suspicious patterns detected.'
    return syntax_note, risk, details