#!/usr/bin/env python3
"""
PowerShell Deobfuscator - Multi-layer recursive deobfuscation of common
PowerShell obfuscation techniques.

Supports: string concatenation, character code conversion, base64 encoding,
invoke-expression layering, variable substitution, backtick insertion,
format string abuse, and more.

Usage:
    python3 powershell_deobfuscator.py --input script.ps1
    python3 powershell_deobfuscator.py --input-string "obfuscated_code"
    python3 powershell_deobfuscator.py --input script.ps1 --max-depth 10 --verbose
"""

from __future__ import annotations

import argparse
import base64
import codecs
import json
import logging
import os
import re
import sys
from pathlib import Path
from typing import List, Optional, Tuple

logger = logging.getLogger("ps_deobfuscator")


class DeobfuscationLayer:
    """Represents a single deobfuscation transformation."""

    def __init__(self, technique: str, before: str, after: str):
        self.technique = technique
        self.before_snippet = before[:200] if len(before) > 200 else before
        self.after_snippet = after[:200] if len(after) > 200 else after

    def to_dict(self) -> dict:
        return {
            "technique": self.technique,
            "before": self.before_snippet,
            "after": self.after_snippet,
        }


class PowerShellDeobfuscator:
    """Multi-layer PowerShell script deobfuscator."""

    def __init__(self, max_depth: int = 15, verbose: bool = False):
        self.max_depth = max_depth
        self.verbose = verbose
        self.layers: List[DeobfuscationLayer] = []

    def deobfuscate(self, script: str) -> str:
        """Apply all deobfuscation techniques iteratively until stable."""
        previous = ""
        depth = 0

        while script != previous and depth < self.max_depth:
            previous = script
            depth += 1

            if self.verbose:
                logger.info(f"=== Deobfuscation pass {depth} ===")

            script = self._remove_backticks(script)
            script = self._resolve_string_concatenation(script)
            script = self._resolve_char_codes(script)
            script = self._resolve_format_strings(script)
            script = self._decode_base64_strings(script)
            script = self._resolve_replace_methods(script)
            script = self._resolve_string_reverse(script)
            script = self._unwrap_invoke_expression(script)
            script = self._decode_encoded_command(script)
            script = self._resolve_chr_join(script)
            script = self._resolve_compressed_streams(script)
            script = self._resolve_bxor(script)
            script = self._clean_whitespace(script)

        if depth >= self.max_depth:
            logger.warning(
                f"Reached maximum deobfuscation depth ({self.max_depth}). "
                "Output may still contain obfuscation."
            )

        return script

    def _record_layer(self, technique: str, before: str, after: str) -> None:
        if before != after:
            self.layers.append(DeobfuscationLayer(technique, before, after))
            if self.verbose:
                logger.info(f"  Applied: {technique}")

    def _remove_backticks(self, script: str) -> str:
        """Remove PowerShell escape backticks used for obfuscation."""
        original = script
        # Remove backticks before non-special characters (obfuscation)
        # Keep backticks for actual escape sequences: `n, `r, `t, `0, `a, `b, `f, `v
        result = re.sub(r'`([^nrt0abfv`\n])', r'\1', script)
        self._record_layer("backtick_removal", original, result)
        return result

    def _resolve_string_concatenation(self, script: str) -> str:
        """Resolve string concatenation: 'abc' + 'def' -> 'abcdef'."""
        original = script
        # Match 'str1' + 'str2' patterns (single and double quotes)
        pattern = r"""(['"])([^'"]*?)\1\s*\+\s*(['"])([^'"]*?)\3"""
        max_iterations = 200
        iteration = 0
        while re.search(pattern, script) and iteration < max_iterations:
            script = re.sub(pattern, lambda m: m.group(1) + m.group(2) + m.group(4) + m.group(1), script, count=1)
            iteration += 1
        self._record_layer("string_concatenation", original, script)
        return script

    def _resolve_char_codes(self, script: str) -> str:
        """Resolve [char]N and [char]0xN patterns."""
        original = script

        # [char]65 or [CHAR]65
        def replace_char_int(m) -> dict:
            try:
                return repr(chr(int(m.group(1))))[1:-1]
            except (ValueError, OverflowError):
                return m.group(0)

        script = re.sub(r'\[char\]\s*(\d+)', replace_char_int, script, flags=re.IGNORECASE)

        # [char]0x41
        def replace_char_hex(m) -> dict:
            try:
                return repr(chr(int(m.group(1), 16)))[1:-1]
            except (ValueError, OverflowError):
                return m.group(0)

        script = re.sub(r'\[char\]\s*0x([0-9a-fA-F]+)', replace_char_hex, script, flags=re.IGNORECASE)

        self._record_layer("char_code_resolution", original, script)
        return script

    def _resolve_format_strings(self, script: str) -> str:
        """Resolve PowerShell format string abuse: '{0}{1}' -f 'A','B'."""
        original = script

        def resolve_format(m) -> dict:
            fmt_string = m.group(1)
            args_str = m.group(2)
            try:
                args = []
                # Parse comma-separated quoted strings
                for arg_match in re.finditer(r"""['"]([^'"]*?)['"]""", args_str):
                    args.append(arg_match.group(1))
                if not args:
                    return m.group(0)
                # Replace {0}, {1}, etc.
                result = fmt_string
                for i, arg in enumerate(args):
                    result = result.replace('{' + str(i) + '}', arg)
                return "'" + result + "'"
            except (IndexError, KeyError):
                return m.group(0)

        pattern = r"""(['"])(\{[\d]+\}[^'"]*?)\1\s*-f\s*((?:['"][^'"]*?['"],?\s*)+)"""
        script = re.sub(pattern, resolve_format, script, flags=re.IGNORECASE)

        self._record_layer("format_string_resolution", original, script)
        return script

    def _decode_base64_strings(self, script: str) -> str:
        """Decode base64-encoded strings found in various patterns."""
        original = script

        # [Convert]::FromBase64String('...')
        def decode_b64(m) -> dict:
            try:
                decoded = base64.b64decode(m.group(1)).decode('utf-8', errors='replace')
                return "'" + decoded + "'"
            except Exception:
                return m.group(0)

        script = re.sub(
            r"""\[(?:System\.)?Convert\]::FromBase64String\(\s*['"]([A-Za-z0-9+/=]+)['"]\s*\)""",
            decode_b64,
            script,
            flags=re.IGNORECASE,
        )

        # [Text.Encoding]::UTF8.GetString([Convert]::FromBase64String('...'))
        def decode_b64_utf8(m) -> dict:
            try:
                decoded = base64.b64decode(m.group(1)).decode('utf-8', errors='replace')
                return "'" + decoded + "'"
            except Exception:
                return m.group(0)

        script = re.sub(
            r"""\[(?:System\.)?Text\.Encoding\]::UTF8\.GetString\(\s*\[(?:System\.)?Convert\]::FromBase64String\(\s*['"]([A-Za-z0-9+/=]+)['"]\s*\)\s*\)""",
            decode_b64_utf8,
            script,
            flags=re.IGNORECASE,
        )

        self._record_layer("base64_decode", original, script)
        return script

    def _resolve_replace_methods(self, script: str) -> str:
        """Resolve .Replace('a','b') and -replace 'a','b' patterns."""
        original = script

        # .Replace('old','new')
        def do_replace(m) -> dict:
            full = m.group(0)
            string_content = m.group(1)
            old_str = m.group(2)
            new_str = m.group(3) if m.group(3) is not None else ""
            try:
                result = string_content.replace(old_str, new_str)
                return "'" + result + "'"
            except Exception:
                return full

        pattern = r"""['"]([^'"]*?)['"]\s*\.Replace\(\s*['"]([^'"]*?)['"],\s*['"]([^'"]*?)?['"]\s*\)"""
        max_iterations = 50
        iteration = 0
        while re.search(pattern, script, re.IGNORECASE) and iteration < max_iterations:
            script = re.sub(pattern, do_replace, script, count=1, flags=re.IGNORECASE)
            iteration += 1

        self._record_layer("replace_resolution", original, script)
        return script

    def _resolve_string_reverse(self, script: str) -> str:
        """Resolve string reversal patterns."""
        original = script

        # Detect reversed strings with explicit array reverse
        def reverse_str(m) -> dict:
            try:
                return "'" + m.group(1)[::-1] + "'"
            except Exception:
                return m.group(0)

        # -join[char[]]'string'[size..0] or similar reverse index
        script = re.sub(
            r"""-join\s*\(\s*['"]([^'"]+)['"]\s*\[-1\.\.\s*-\d+\]\s*\)""",
            reverse_str, script, flags=re.IGNORECASE
        )

        self._record_layer("string_reverse", original, script)
        return script

    def _unwrap_invoke_expression(self, script: str) -> str:
        """Remove IEX/Invoke-Expression wrappers to reveal inner content."""
        original = script

        # Remove IEX( ... ) wrapping - keep the inner content
        iex_patterns = [
            r'(?:IEX|Invoke-Expression)\s*\(\s*(.*)\s*\)\s*$',
            r'(?:IEX|Invoke-Expression)\s*\(\s*(.*)\s*\)',
            r'(?:IEX|Invoke-Expression)\s+(.*)',
            r'&\s*\(\s*[\'"]IEX[\'"]\s*\)\s*\(\s*(.*)\s*\)',
            r'\.\s*\(\s*[\'"]IEX[\'"]\s*\)\s*(.*)',
        ]

        for pattern in iex_patterns:
            match = re.search(pattern, script, re.IGNORECASE | re.DOTALL)
            if match:
                inner = match.group(1).strip()
                if inner and len(inner) > 10:
                    script = inner
                    break

        self._record_layer("iex_unwrap", original, script)
        return script

    def _decode_encoded_command(self, script: str) -> str:
        """Decode PowerShell -EncodedCommand / -enc base64 payloads."""
        original = script

        def decode_enc(m) -> dict:
            try:
                decoded = base64.b64decode(m.group(1)).decode('utf-16-le', errors='replace')
                return decoded
            except Exception:
                return m.group(0)

        script = re.sub(
            r'(?:powershell(?:\.exe)?)\s+.*?(?:-[Ee](?:nc(?:odedcommand)?)?)\s+([A-Za-z0-9+/=]{20,})',
            decode_enc,
            script,
            flags=re.IGNORECASE,
        )

        self._record_layer("encoded_command_decode", original, script)
        return script

    def _resolve_chr_join(self, script: str) -> str:
        """Resolve patterns like [string]::join('', (72,101,108) | %{[char]$_})."""
        original = script

        # Find patterns with arrays of integers being converted to chars
        def resolve_join(m) -> dict:
            try:
                numbers = re.findall(r'\d+', m.group(1))
                result = ''.join(chr(int(n)) for n in numbers)
                return "'" + result + "'"
            except (ValueError, OverflowError):
                return m.group(0)

        # [string]::join('', (N,N,N,...) | %{[char]$_})
        script = re.sub(
            r"""\[string\]::join\(\s*['"]['"],\s*\(\s*([\d,\s]+)\s*\)\s*\|\s*%\s*\{\s*\[char\]\s*\$_\s*\}\s*\)""",
            resolve_join, script, flags=re.IGNORECASE,
        )

        # -join(N,N,N | %{[char]$_})
        script = re.sub(
            r"""-join\s*\(\s*\(?\s*([\d,\s]+)\s*\)?\s*\|\s*%\s*\{\s*\[char\]\s*\$_\s*\}\s*\)""",
            resolve_join, script, flags=re.IGNORECASE,
        )

        # (N,N,N) | ForEach-Object {[char]$_}) -join ''
        script = re.sub(
            r"""\(\s*([\d,\s]+)\s*\)\s*\|\s*(?:ForEach-Object|%)\s*\{\s*\[char\]\s*\$_\s*\}\s*(?:\)\s*)?-join\s*['"]['"]\s*""",
            resolve_join, script, flags=re.IGNORECASE,
        )

        self._record_layer("chr_join_resolution", original, script)
        return script

    def _resolve_compressed_streams(self, script: str) -> str:
        """Detect and note compressed/deflate stream patterns (cannot fully decompress without execution)."""
        original = script

        # Flag compressed streams
        compress_patterns = [
            r'IO\.Compression\.DeflateStream',
            r'IO\.Compression\.GZipStream',
            r'New-Object\s+IO\.MemoryStream',
        ]

        for pattern in compress_patterns:
            if re.search(pattern, script, re.IGNORECASE):
                if self.verbose:
                    logger.info("  Detected compressed stream - may require dynamic execution to fully decode")
                break

        return script

    def _resolve_bxor(self, script: str) -> str:
        """Resolve simple XOR obfuscation patterns."""
        original = script

        # Pattern: [byte[]]@(N,N,N) | %{ $_ -bxor KEY }
        def resolve_xor(m) -> dict:
            try:
                numbers = [int(n.strip()) for n in m.group(1).split(',')]
                key = int(m.group(2))
                result = ''.join(chr(b ^ key) for b in numbers)
                return "'" + result + "'"
            except (ValueError, OverflowError):
                return m.group(0)

        script = re.sub(
            r"""\[byte\[\]\]\s*@?\(\s*([\d,\s]+)\s*\)\s*\|\s*%\s*\{\s*\$_\s*-bxor\s+(\d+)\s*\}""",
            resolve_xor, script, flags=re.IGNORECASE,
        )

        self._record_layer("bxor_resolution", original, script)
        return script

    def _clean_whitespace(self, script: str) -> str:
        """Normalize excessive whitespace."""
        original = script
        script = re.sub(r'[ \t]{4,}', '  ', script)
        script = re.sub(r'\n{3,}', '\n\n', script)
        return script

    def get_report(self) -> dict:
        """Generate deobfuscation report."""
        techniques_used = list(set(layer.technique for layer in self.layers))
        return {
            "total_layers": len(self.layers),
            "techniques_used": techniques_used,
            "layers": [layer.to_dict() for layer in self.layers],
        }


def main() -> None:
    parser = argparse.ArgumentParser(
        description="PowerShell script deobfuscator - recursively removes obfuscation layers",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s --input encoded_script.ps1
  %(prog)s --input-string "IEX([string]::join('',((87,104) | %%{[char]$_})))"
  %(prog)s --input script.ps1 --max-depth 10 --verbose --output decoded.ps1
  %(prog)s --input script.ps1 --report report.json
        """,
    )
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument("--input", "-i", help="Path to obfuscated PowerShell script")
    input_group.add_argument("--input-string", "-s", help="Obfuscated PowerShell string")
    parser.add_argument("--output", "-o", help="Output file for deobfuscated script")
    parser.add_argument("--report", "-r", help="Output file for JSON deobfuscation report")
    parser.add_argument("--max-depth", type=int, default=15,
                        help="Maximum deobfuscation recursion depth (default: 15)")
    parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose output")
    parser.add_argument("--format", default="json", choices=["json", "text", "csv"],
                        help="Output format (default: json)")

    args = parser.parse_args()

    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=log_level, format="%(levelname)s: %(message)s")

    # Read input
    if args.input:
        input_path = os.path.abspath(args.input)
        if not os.path.isfile(input_path):
            logger.error(f"Input file not found: {input_path}")
            sys.exit(1)
        try:
            with open(input_path, "r", encoding="utf-8", errors="replace") as f:
                script = f.read()
        except OSError as e:
            logger.error(f"Failed to read input file: {e}")
            sys.exit(1)
    else:
        script = args.input_string

    logger.info(f"Input length: {len(script)} characters")

    # Deobfuscate
    deobfuscator = PowerShellDeobfuscator(max_depth=args.max_depth, verbose=args.verbose)
    result = deobfuscator.deobfuscate(script)

    # Output result
    if args.output:
        output_path = os.path.abspath(args.output)
        try:
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(result)
            logger.info(f"Deobfuscated output written to: {output_path}")
        except OSError as e:
            logger.error(f"Failed to write output: {e}")
            sys.exit(1)
    else:
        print("\n" + "=" * 60)
        print("DEOBFUSCATED OUTPUT:")
        print("=" * 60)
        print(result)
        print("=" * 60)

    # Report
    report = deobfuscator.get_report()
    logger.info(f"Total deobfuscation layers applied: {report['total_layers']}")
    if report['techniques_used']:
        logger.info(f"Techniques found: {', '.join(report['techniques_used'])}")

    if args.report:
        report_path = os.path.abspath(args.report)
        try:
            with open(report_path, "w", encoding="utf-8") as f:
                json.dump(report, f, indent=2)
            logger.info(f"Report written to: {report_path}")
        except OSError as e:
            logger.error(f"Failed to write report: {e}")
            sys.exit(1)


if __name__ == "__main__":
    main()