#!/usr/bin/env python3
"""PowerShell obfuscated malware deobfuscation agent."""

import json
import argparse
import re
import base64
from datetime import datetime


def decode_base64_commands(script_content):
    """Find and decode Base64 encoded PowerShell commands."""
    decoded = []
    b64_pattern = re.compile(r'-[eE](?:nc(?:odedcommand)?)\s+([A-Za-z0-9+/=]{20,})')
    for match in b64_pattern.finditer(script_content):
        encoded = match.group(1)
        try:
            raw = base64.b64decode(encoded)
            text = raw.decode("utf-16-le", errors="replace")
            decoded.append({"encoded": encoded[:60] + "...", "decoded": text[:500]})
        except Exception:
            pass
    standalone_b64 = re.compile(r'["\']([A-Za-z0-9+/]{40,}={0,2})["\']')
    for match in standalone_b64.finditer(script_content):
        try:
            raw = base64.b64decode(match.group(1))
            text = raw.decode("utf-8", errors="replace")
            if text.isprintable() or "http" in text.lower():
                decoded.append({"encoded": match.group(1)[:60] + "...", "decoded": text[:500]})
        except Exception:
            pass
    return decoded


def deobfuscate_string_concatenation(script_content):
    """Resolve string concatenation obfuscation."""
    concat_pattern = re.compile(r"(?:'[^']*'\s*\+\s*){2,}'[^']*'")
    resolved = []
    for match in concat_pattern.finditer(script_content):
        original = match.group(0)
        parts = re.findall(r"'([^']*)'", original)
        result = "".join(parts)
        resolved.append({"obfuscated": original[:80], "resolved": result[:500]})
    return resolved


def detect_obfuscation_techniques(script_content):
    """Identify obfuscation techniques used in the script."""
    techniques = []
    checks = [
        (r'-[eE](?:nc(?:odedcommand)?)', "Base64 encoded command", "HIGH"),
        (r'\[(?:char|int)\]\s*\d+', "Character code conversion", "MEDIUM"),
        (r'(?:iex|invoke-expression)', "Invoke-Expression (IEX) execution", "HIGH"),
        (r'\$\{[^}]+\}', "Variable name obfuscation with braces", "LOW"),
        (r'\.(?:replace|split|reverse)\(', "String manipulation methods", "MEDIUM"),
        (r'-(?:join|split)\s', "Array join/split obfuscation", "MEDIUM"),
        (r'(?:Net\.WebClient|DownloadString|DownloadFile)', "Web download cradle", "CRITICAL"),
        (r'(?:Start-Process|Invoke-Item|cmd\s*/c)', "Process execution", "HIGH"),
        (r'\[System\.Convert\]::FromBase64String', ".NET Base64 decode", "HIGH"),
        (r'(?:gci|ls|dir)\s+env:', "Environment variable access", "LOW"),
    ]
    for pattern, name, severity in checks:
        if re.search(pattern, script_content, re.IGNORECASE):
            techniques.append({"technique": name, "severity": severity})
    return techniques


def extract_iocs(script_content):
    """Extract indicators of compromise from deobfuscated content."""
    iocs = {"urls": [], "ips": [], "domains": [], "file_paths": []}
    url_pattern = re.compile(r'https?://[^\s"\'<>]+')
    ip_pattern = re.compile(r'\b(?:\d{1,3}\.){3}\d{1,3}\b')
    path_pattern = re.compile(r'[A-Z]:\\[\w\\]+\.\w{2,4}|/(?:tmp|var|etc)/[\w/]+')
    iocs["urls"] = list(set(url_pattern.findall(script_content)))
    iocs["ips"] = list(set(ip_pattern.findall(script_content)))
    iocs["file_paths"] = list(set(path_pattern.findall(script_content)))
    return iocs


def run_analysis(script_path):
    """Execute PowerShell deobfuscation analysis."""
    print(f"\n{'='*60}")
    print(f"  POWERSHELL MALWARE DEOBFUSCATION")
    print(f"  File: {script_path}")
    print(f"  Generated: {datetime.utcnow().isoformat()} UTC")
    print(f"{'='*60}\n")

    with open(script_path, "r", errors="replace") as f:
        content = f.read()

    techniques = detect_obfuscation_techniques(content)
    print(f"--- OBFUSCATION TECHNIQUES ({len(techniques)}) ---")
    for t in techniques:
        print(f"  [{t['severity']}] {t['technique']}")

    b64 = decode_base64_commands(content)
    print(f"\n--- BASE64 DECODED ({len(b64)}) ---")
    for d in b64[:5]:
        print(f"  {d['decoded'][:100]}")

    concat = deobfuscate_string_concatenation(content)
    print(f"\n--- STRING CONCAT RESOLVED ({len(concat)}) ---")
    for c in concat[:5]:
        print(f"  {c['resolved'][:100]}")

    all_decoded = content
    for d in b64:
        all_decoded += "\n" + d["decoded"]
    iocs = extract_iocs(all_decoded)
    print(f"\n--- IOCs ---")
    print(f"  URLs: {iocs['urls'][:5]}")
    print(f"  IPs: {iocs['ips'][:5]}")
    print(f"  Paths: {iocs['file_paths'][:5]}")

    return {"techniques": techniques, "decoded_b64": b64, "concat": concat, "iocs": iocs}


def main():
    parser = argparse.ArgumentParser(description="PowerShell Deobfuscation Agent")
    parser.add_argument("--script", required=True, help="Path to obfuscated PowerShell script")
    parser.add_argument("--output", help="Save report to JSON file")
    args = parser.parse_args()

    report = run_analysis(args.script)
    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2, default=str)
        print(f"\n[+] Report saved to {args.output}")


if __name__ == "__main__":
    main()