#!/usr/bin/env python3
"""
APktool MCP Server - Gradio Space
Expõe funcionalidades do Apktool como MCP server para HuggingChat
"""

import gradio as gr
import subprocess
import tempfile
import shutil
import os
from pathlib import Path
import json
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Diretório de trabalho persistente
WORK_DIR = Path("/tmp/apktool_workspace")
WORK_DIR.mkdir(exist_ok=True)


def decode_apk(apk_file, output_name: str = "", force: bool = False, no_res: bool = False, no_src: bool = False) -> str:
    """
    Decompile an APK file to extract resources, manifest, and smali code.
    
    Args:
        apk_file: The APK file to decompile
        output_name: Output directory name (optional, defaults to APK name)
        force: Force overwrite existing directory
        no_res: Do not decode resources
        no_src: Do not decode sources (smali)
    
    Returns:
        JSON string with decompiled directory path and status
    """
    if apk_file is None:
        return json.dumps({"error": "No APK file provided"})
    
    try:
        apk_path = Path(apk_file)
        
        # Determinar diretório de saída
        if output_name:
            output_dir = WORK_DIR / output_name
        else:
            output_dir = WORK_DIR / apk_path.stem
        
        # Construir comando
        cmd = ["apktool", "d", str(apk_path), "-o", str(output_dir)]
        if force:
            cmd.append("-f")
        if no_res:
            cmd.append("-r")
        if no_src:
            cmd.append("-s")
        
        # Executar
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        
        if result.returncode == 0:
            return json.dumps({
                "success": True,
                "output_dir": str(output_dir),
                "message": f"APK decompiled successfully to {output_dir}",
                "stdout": result.stdout
            })
        else:
            return json.dumps({
                "success": False,
                "error": result.stderr,
                "message": "Failed to decompile APK"
            })
    except subprocess.TimeoutExpired:
        return json.dumps({"success": False, "error": "Operation timed out"})
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def build_apk(source_dir: str, output_path: str = "") -> str:
    """
    Recompile/build an APK from decompiled source directory.
    
    Args:
        source_dir: Path to the decompiled source directory
        output_path: Output APK path (optional)
    
    Returns:
        JSON string with built APK path and status
    """
    try:
        src_path = WORK_DIR / source_dir if not source_dir.startswith("/") else Path(source_dir)
        
        if not src_path.exists():
            return json.dumps({"success": False, "error": f"Source directory not found: {src_path}"})
        
        cmd = ["apktool", "b", str(src_path)]
        if output_path:
            cmd.extend(["-o", output_path])
        
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        
        if result.returncode == 0:
            built_apk = src_path / "dist" / f"{src_path.name}.apk"
            return json.dumps({
                "success": True,
                "apk_path": str(built_apk) if built_apk.exists() else "Unknown",
                "message": "APK built successfully",
                "stdout": result.stdout
            })
        else:
            return json.dumps({
                "success": False,
                "error": result.stderr,
                "message": "Failed to build APK"
            })
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def analyze_manifest(apk_name: str) -> str:
    """
    Parse AndroidManifest.xml for permissions and components.
    
    Args:
        apk_name: Name of the decompiled APK directory
    
    Returns:
        JSON string with manifest analysis
    """
    try:
        manifest_path = WORK_DIR / apk_name / "AndroidManifest.xml"
        
        if not manifest_path.exists():
            return json.dumps({"success": False, "error": f"Manifest not found. Decompile APK first."})
        
        with open(manifest_path, 'r') as f:
            content = f.read()
        
        # Extrair informações básicas
        import re
        
        package = re.search(r'package="([^"]+)"', content)
        permissions = re.findall(r'<uses-permission[^>]*android:name="([^"]+)"', content)
        activities = re.findall(r'<activity[^>]*android:name="([^"]+)"', content)
        services = re.findall(r'<service[^>]*android:name="([^"]+)"', content)
        receivers = re.findall(r'<receiver[^>]*android:name="([^"]+)"', content)
        
        return json.dumps({
            "success": True,
            "package": package.group(1) if package else "Unknown",
            "permissions": permissions,
            "activities": activities,
            "services": services,
            "receivers": receivers,
            "permission_count": len(permissions),
            "component_count": len(activities) + len(services) + len(receivers)
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def list_permissions(apk_name: str) -> str:
    """
    List all permissions requested by an APK.
    
    Args:
        apk_name: Name of the decompiled APK directory
    
    Returns:
        JSON string with permissions list
    """
    try:
        manifest_path = WORK_DIR / apk_name / "AndroidManifest.xml"
        
        if not manifest_path.exists():
            return json.dumps({"success": False, "error": "Manifest not found. Decompile APK first."})
        
        with open(manifest_path, 'r') as f:
            content = f.read()
        
        import re
        permissions = re.findall(r'<uses-permission[^>]*android:name="([^"]+)"', content)
        
        # Categorizar permissões perigosas
        dangerous_permissions = [
            "READ_CONTACTS", "WRITE_CONTACTS", "READ_CALENDAR", "WRITE_CALENDAR",
            "CAMERA", "READ_CALL_LOG", "WRITE_CALL_LOG", "PROCESS_OUTGOING_CALLS",
            "ACCESS_FINE_LOCATION", "ACCESS_COARSE_LOCATION", "RECORD_AUDIO",
            "READ_PHONE_STATE", "CALL_PHONE", "ANSWER_PHONE_CALLS",
            "READ_EXTERNAL_STORAGE", "WRITE_EXTERNAL_STORAGE", "SEND_SMS", "RECEIVE_SMS"
        ]
        
        dangerous = []
        normal = []
        
        for p in permissions:
            is_dangerous = any(d in p.upper() for d in dangerous_permissions)
            if is_dangerous:
                dangerous.append(p)
            else:
                normal.append(p)
        
        return json.dumps({
            "success": True,
            "total_permissions": len(permissions),
            "dangerous_permissions": dangerous,
            "normal_permissions": normal,
            "dangerous_count": len(dangerous)
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def extract_strings(apk_name: str, locale: str = "") -> str:
    """
    Extract string resources from a decompiled APK.
    
    Args:
        apk_name: Name of the decompiled APK directory
        locale: Locale to extract (e.g., 'en', 'pt'). Empty for default.
    
    Returns:
        JSON string with extracted strings
    """
    try:
        if locale:
            strings_path = WORK_DIR / apk_name / "res" / f"values-{locale}" / "strings.xml"
        else:
            strings_path = WORK_DIR / apk_name / "res" / "values" / "strings.xml"
        
        if not strings_path.exists():
            return json.dumps({"success": False, "error": f"Strings file not found at {strings_path}"})
        
        with open(strings_path, 'r') as f:
            content = f.read()
        
        import re
        strings = re.findall(r'<string name="([^"]+)"[^>]*>([^<]*)</string>', content)
        
        return json.dumps({
            "success": True,
            "file_path": str(strings_path),
            "string_count": len(strings),
            "strings": {name: value for name, value in strings[:50]}  # Limitar a 50
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def find_smali_references(apk_name: str, pattern: str) -> str:
    """
    Search for patterns in decompiled smali code.
    
    Args:
        apk_name: Name of the decompiled APK directory
        pattern: Search pattern (regex supported)
    
    Returns:
        JSON string with search results
    """
    try:
        smali_dir = WORK_DIR / apk_name / "smali"
        
        if not smali_dir.exists():
            return json.dumps({"success": False, "error": "Smali directory not found. Decompile APK first."})
        
        import re
        results = []
        regex = re.compile(pattern, re.IGNORECASE)
        
        for smali_file in smali_dir.rglob("*.smali"):
            try:
                with open(smali_file, 'r', errors='ignore') as f:
                    for i, line in enumerate(f, 1):
                        if regex.search(line):
                            results.append({
                                "file": str(smali_file.relative_to(smali_dir)),
                                "line": i,
                                "content": line.strip()[:200]
                            })
                            if len(results) >= 100:  # Limitar resultados
                                break
            except:
                continue
            if len(results) >= 100:
                break
        
        return json.dumps({
            "success": True,
            "pattern": pattern,
            "matches": len(results),
            "results": results
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def get_apk_info(apk_file) -> str:
    """
    Get basic APK metadata and information.
    
    Args:
        apk_file: The APK file to analyze
    
    Returns:
        JSON string with APK information
    """
    if apk_file is None:
        return json.dumps({"error": "No APK file provided"})
    
    try:
        apk_path = Path(apk_file)
        
        # Usar aapt para info básica
        result = subprocess.run(
            ["aapt", "dump", "badging", str(apk_path)],
            capture_output=True, text=True, timeout=60
        )
        
        info = {
            "success": True,
            "file_name": apk_path.name,
            "file_size_mb": round(apk_path.stat().st_size / (1024 * 1024), 2),
        }
        
        if result.returncode == 0:
            import re
            package = re.search(r"package: name='([^']+)'", result.stdout)
            version = re.search(r"versionName='([^']+)'", result.stdout)
            sdk = re.search(r"targetSdkVersion:'([^']+)'", result.stdout)
            
            if package:
                info["package_name"] = package.group(1)
            if version:
                info["version"] = version.group(1)
            if sdk:
                info["target_sdk"] = sdk.group(1)
        
        return json.dumps(info, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def list_decompiled() -> str:
    """
    List all decompiled APKs in the workspace.
    
    Returns:
        JSON string with list of decompiled APKs
    """
    try:
        decompiled = []
        for item in WORK_DIR.iterdir():
            if item.is_dir() and (item / "AndroidManifest.xml").exists():
                decompiled.append(item.name)
        
        return json.dumps({
            "success": True,
            "workspace": str(WORK_DIR),
            "decompiled_apks": decompiled,
            "count": len(decompiled)
        }, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


# Criar a interface Gradio
with gr.Blocks(title="APktool MCP Server") as demo:
    gr.Markdown("# 🔧 APktool MCP Server")
    gr.Markdown("Android APK analysis tools exposed as MCP for HuggingChat")
    
    with gr.Tab("Decode APK"):
        decode_input = gr.File(label="APK File", file_types=[".apk"])
        decode_name = gr.Textbox(label="Output Directory Name (optional)")
        with gr.Row():
            decode_force = gr.Checkbox(label="Force Overwrite", value=False)
            decode_no_res = gr.Checkbox(label="Skip Resources", value=False)
            decode_no_src = gr.Checkbox(label="Skip Sources", value=False)
        decode_btn = gr.Button("Decode", variant="primary")
        decode_output = gr.Textbox(label="Result", lines=10)
        decode_btn.click(decode_apk, [decode_input, decode_name, decode_force, decode_no_res, decode_no_src], decode_output)
    
    with gr.Tab("Build APK"):
        build_dir = gr.Textbox(label="Source Directory Name")
        build_output = gr.Textbox(label="Output APK Path (optional)")
        build_btn = gr.Button("Build", variant="primary")
        build_result = gr.Textbox(label="Result", lines=10)
        build_btn.click(build_apk, [build_dir, build_output], build_result)
    
    with gr.Tab("Analyze"):
        analyze_apk = gr.Textbox(label="Decompiled APK Name")
        analyze_btn = gr.Button("Analyze Manifest", variant="primary")
        analyze_output = gr.Textbox(label="Result", lines=15)
        analyze_btn.click(analyze_manifest, [analyze_apk], analyze_output)
    
    with gr.Tab("Permissions"):
        perm_apk = gr.Textbox(label="Decompiled APK Name")
        perm_btn = gr.Button("List Permissions", variant="primary")
        perm_output = gr.Textbox(label="Result", lines=15)
        perm_btn.click(list_permissions, [perm_apk], perm_output)
    
    with gr.Tab("Strings"):
        strings_apk = gr.Textbox(label="Decompiled APK Name")
        strings_locale = gr.Textbox(label="Locale (e.g., 'en', 'pt')", value="")
        strings_btn = gr.Button("Extract Strings", variant="primary")
        strings_output = gr.Textbox(label="Result", lines=15)
        strings_btn.click(extract_strings, [strings_apk, strings_locale], strings_output)
    
    with gr.Tab("Search Smali"):
        search_apk = gr.Textbox(label="Decompiled APK Name")
        search_pattern = gr.Textbox(label="Search Pattern (regex)", placeholder="e.g., api_key, password, http")
        search_btn = gr.Button("Search", variant="primary")
        search_output = gr.Textbox(label="Result", lines=15)
        search_btn.click(find_smali_references, [search_apk, search_pattern], search_output)
    
    with gr.Tab("APK Info"):
        info_file = gr.File(label="APK File", file_types=[".apk"])
        info_btn = gr.Button("Get Info", variant="primary")
        info_output = gr.Textbox(label="Result", lines=10)
        info_btn.click(get_apk_info, [info_file], info_output)
    
    with gr.Tab("Workspace"):
        list_btn = gr.Button("List Decompiled APKs", variant="primary")
        list_output = gr.Textbox(label="Workspace Contents", lines=10)
        list_btn.click(list_decompiled, [], list_output)


# Lançar como MCP server!
if __name__ == "__main__":
    demo.launch(mcp_server=True)