#!/usr/bin/env python3
"""Agent for detecting insecure data storage in mobile applications."""

import argparse
import json
import os
import re
import subprocess
import sqlite3
from datetime import datetime, timezone


ANDROID_SENSITIVE_PATHS = [
    "/data/data/{package}/shared_prefs/",
    "/data/data/{package}/databases/",
    "/data/data/{package}/files/",
    "/data/data/{package}/cache/",
    "/sdcard/Android/data/{package}/",
]

_SAFE_TABLE_RE = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]*$')

SENSITIVE_PATTERNS = {
    "api_key": re.compile(r'["\']?api[_-]?key["\']?\s*[:=]\s*["\']([^"\']+)', re.I),
    "token": re.compile(r'["\']?(?:access|auth|bearer)[_-]?token["\']?\s*[:=]\s*["\']([^"\']+)', re.I),
    "password": re.compile(r'["\']?password["\']?\s*[:=]\s*["\']([^"\']+)', re.I),
    "private_key": re.compile(r'-----BEGIN (?:RSA )?PRIVATE KEY-----'),
    "base64_cred": re.compile(r'["\']?(?:auth|credential)["\']?\s*[:=]\s*["\']([A-Za-z0-9+/=]{20,})', re.I),
}


def scan_shared_prefs(prefs_dir):
    """Scan Android SharedPreferences XML files for sensitive data."""
    findings = []
    if not os.path.isdir(prefs_dir):
        return findings
    for fname in os.listdir(prefs_dir):
        if not fname.endswith(".xml"):
            continue
        fpath = os.path.join(prefs_dir, fname)
        try:
            with open(fpath, "r", errors="replace") as f:
                content = f.read()
            for pattern_name, pattern in SENSITIVE_PATTERNS.items():
                matches = pattern.findall(content)
                if matches:
                    findings.append({
                        "file": fpath,
                        "type": "shared_prefs",
                        "pattern": pattern_name,
                        "match_count": len(matches),
                        "severity": "HIGH",
                    })
        except PermissionError:
            pass
    return findings


def scan_sqlite_databases(db_dir):
    """Scan SQLite databases for unencrypted sensitive data."""
    findings = []
    if not os.path.isdir(db_dir):
        return findings
    for fname in os.listdir(db_dir):
        if not fname.endswith((".db", ".sqlite", ".sqlite3")):
            continue
        fpath = os.path.join(db_dir, fname)
        try:
            conn = sqlite3.connect(fpath)
            cursor = conn.cursor()
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
            tables = cursor.fetchall()
            for (table_name,) in tables:
                if not _SAFE_TABLE_RE.match(table_name):
                    continue
                cursor.execute(f"PRAGMA table_info([{table_name}])")
                columns = cursor.fetchall()
                sensitive_cols = []
                for col in columns:
                    col_name = col[1].lower()
                    for sf in ["password", "token", "secret", "key", "ssn", "credit"]:
                        if sf in col_name:
                            sensitive_cols.append(col[1])
                if sensitive_cols:
                    cursor.execute(f"SELECT COUNT(*) FROM [{table_name}]")
                    row_count = cursor.fetchone()[0]
                    findings.append({
                        "file": fpath,
                        "table": table_name,
                        "sensitive_columns": sensitive_cols,
                        "row_count": row_count,
                        "encrypted": False,
                        "severity": "CRITICAL",
                    })
            conn.close()
        except (sqlite3.Error, PermissionError):
            pass
    return findings


def scan_file_storage(files_dir):
    """Scan app file storage for sensitive data."""
    findings = []
    if not os.path.isdir(files_dir):
        return findings
    for root, _, files in os.walk(files_dir):
        for fname in files:
            fpath = os.path.join(root, fname)
            try:
                with open(fpath, "r", errors="replace") as f:
                    content = f.read(4096)
                for pattern_name, pattern in SENSITIVE_PATTERNS.items():
                    if pattern.search(content):
                        findings.append({
                            "file": fpath,
                            "pattern": pattern_name,
                            "severity": "HIGH",
                        })
            except (PermissionError, UnicodeDecodeError):
                pass
    return findings


def adb_pull_app_data(package_name, output_dir):
    """Pull application data via ADB for analysis."""
    os.makedirs(output_dir, exist_ok=True)
    paths = [p.format(package=package_name) for p in ANDROID_SENSITIVE_PATHS]
    results = []
    for path in paths:
        try:
            subprocess.check_output(
                ["adb", "pull", path, output_dir],
                text=True, errors="replace", timeout=15
            )
            results.append({"path": path, "status": "pulled"})
        except subprocess.SubprocessError:
            results.append({"path": path, "status": "failed"})
    return results


def main():
    parser = argparse.ArgumentParser(
        description="Detect insecure data storage in mobile apps (authorized testing only)"
    )
    parser.add_argument("--scan-dir", help="Directory containing app data to scan")
    parser.add_argument("--package", help="Android package name for ADB pull")
    parser.add_argument("--pull-dir", default=os.environ.get("MOBILE_AUDIT_DIR", "/tmp/mobile_audit"))
    parser.add_argument("--output", "-o", help="Output JSON report")
    args = parser.parse_args()

    print("[*] Insecure Mobile Data Storage Detection Agent")
    report = {"timestamp": datetime.now(timezone.utc).isoformat(), "findings": []}

    scan_dir = args.scan_dir or args.pull_dir
    if args.package:
        adb_pull_app_data(args.package, args.pull_dir)

    if os.path.isdir(scan_dir):
        report["findings"].extend(scan_shared_prefs(scan_dir))
        report["findings"].extend(scan_sqlite_databases(scan_dir))
        report["findings"].extend(scan_file_storage(scan_dir))

    critical = sum(1 for f in report["findings"] if f.get("severity") == "CRITICAL")
    report["risk_level"] = "CRITICAL" if critical else "HIGH" if report["findings"] else "LOW"
    print(f"[*] Findings: {len(report['findings'])} (CRITICAL: {critical})")

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"[*] Report saved to {args.output}")
    else:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()
