#!/usr/bin/env python3
"""Go malware analysis agent for Ghidra-assisted reverse engineering.

Analyzes Go binaries to extract function names, strings, build metadata,
package information, and detects common Go malware characteristics.
"""

import os
import sys
import json
import hashlib
import re
import math
from collections import Counter


def compute_hash(filepath):
    """Compute SHA-256 hash of file."""
    sha256 = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            sha256.update(chunk)
    return sha256.hexdigest()


def shannon_entropy(data):
    """Calculate Shannon entropy."""
    if not data:
        return 0.0
    freq = Counter(data)
    length = len(data)
    return -sum((c / length) * math.log2(c / length) for c in freq.values())


def detect_go_binary(filepath):
    """Detect if a binary is compiled with Go and extract version info."""
    with open(filepath, "rb") as f:
        data = f.read()

    indicators = {
        "is_go_binary": False,
        "go_version": None,
        "go_buildinfo": False,
        "gopclntab_found": False,
    }

    # Go build info magic
    buildinfo_magic = b"\xff Go buildinf:"
    offset = data.find(buildinfo_magic)
    if offset != -1:
        indicators["is_go_binary"] = True
        indicators["go_buildinfo"] = True

    # Go version string
    version_pattern = rb"go(\d+\.\d+(?:\.\d+)?)"
    matches = re.findall(version_pattern, data)
    if matches:
        indicators["is_go_binary"] = True
        versions = sorted(set(m.decode() for m in matches))
        indicators["go_version"] = versions[-1] if versions else None

    # gopclntab (Go PC line table) magic bytes
    gopclntab_magics = [
        b"\xfb\xff\xff\xff\x00\x00",  # Go 1.2-1.15
        b"\xfa\xff\xff\xff\x00\x00",  # Go 1.16-1.17
        b"\xf0\xff\xff\xff\x00\x00",  # Go 1.18+
        b"\xf1\xff\xff\xff\x00\x00",  # Go 1.20+
    ]
    for magic in gopclntab_magics:
        if magic in data:
            indicators["gopclntab_found"] = True
            indicators["is_go_binary"] = True
            break

    # Runtime strings
    go_strings = [b"runtime.main", b"runtime.goexit", b"runtime.gopanic",
                  b"runtime.newproc", b"GOROOT", b"GOPATH"]
    found_runtime = sum(1 for s in go_strings if s in data)
    if found_runtime >= 2:
        indicators["is_go_binary"] = True
    indicators["runtime_strings_found"] = found_runtime

    return indicators


def extract_go_strings(filepath, min_length=6):
    """Extract Go-style strings (length-prefixed, not null-terminated)."""
    with open(filepath, "rb") as f:
        data = f.read()

    # Standard ASCII string extraction
    ascii_pattern = re.compile(rb"[\x20-\x7e]{%d,}" % min_length)
    strings = [m.group().decode("ascii", errors="replace") for m in ascii_pattern.finditer(data)]
    return strings


def extract_go_packages(strings_list):
    """Identify Go packages from extracted strings."""
    packages = set()
    pkg_pattern = re.compile(r"^([a-zA-Z0-9_]+(?:/[a-zA-Z0-9_.-]+)+)\.")
    for s in strings_list:
        match = pkg_pattern.match(s)
        if match:
            packages.add(match.group(1))
    # Also look for known Go import paths
    for s in strings_list:
        if s.startswith("github.com/") or s.startswith("golang.org/"):
            parts = s.split("/")
            if len(parts) >= 3:
                packages.add("/".join(parts[:3]))
    return sorted(packages)


SUSPICIOUS_GO_PACKAGES = {
    "github.com/kbinani/screenshot": "Screen capture capability",
    "github.com/atotto/clipboard": "Clipboard access",
    "github.com/go-vgo/robotgo": "Desktop automation / keylogging",
    "github.com/miekg/dns": "Custom DNS resolution (C2/tunneling)",
    "golang.org/x/crypto/ssh": "SSH client (lateral movement)",
    "github.com/shirou/gopsutil": "System enumeration",
    "github.com/mitchellh/go-ps": "Process listing",
    "github.com/gobuffalo/packr": "Binary resource embedding",
    "github.com/Ne0nd0g/merlin": "Merlin C2 agent",
    "github.com/BishopFox/sliver": "Sliver C2 framework",
    "github.com/traefik/yaegi": "Go interpreter (dynamic execution)",
}


def detect_suspicious_packages(packages):
    """Flag suspicious Go packages commonly used in malware."""
    findings = []
    for pkg in packages:
        for sus_pkg, description in SUSPICIOUS_GO_PACKAGES.items():
            if sus_pkg in pkg:
                findings.append({"package": pkg, "concern": description})
    return findings


def analyze_sections(filepath):
    """Analyze PE/ELF sections for Go binary characteristics."""
    with open(filepath, "rb") as f:
        magic = f.read(4)
        f.seek(0)
        data = f.read()

    sections = []
    if magic[:2] == b"MZ":  # PE
        try:
            import pefile
            pe = pefile.PE(data=data)
            for section in pe.sections:
                name = section.Name.rstrip(b"\x00").decode("ascii", errors="replace")
                entropy = section.get_entropy()
                sections.append({
                    "name": name, "virtual_size": section.Misc_VirtualSize,
                    "raw_size": section.SizeOfRawData, "entropy": round(entropy, 3),
                })
            pe.close()
        except ImportError:
            sections.append({"note": "pefile not installed"})
    elif magic[:4] == b"\x7fELF":
        try:
            from elftools.elf.elffile import ELFFile
            from io import BytesIO
            elf = ELFFile(BytesIO(data))
            for section in elf.iter_sections():
                sec_data = section.data() if section.header.sh_size > 0 else b""
                entropy = shannon_entropy(sec_data) if sec_data else 0
                sections.append({
                    "name": section.name, "size": section.header.sh_size,
                    "entropy": round(entropy, 3), "type": section.header.sh_type,
                })
        except ImportError:
            sections.append({"note": "pyelftools not installed"})
    return sections


def detect_obfuscation(go_info, strings_list):
    """Detect Go binary obfuscation (garble, gobfuscate)."""
    indicators = {"obfuscated": False, "techniques": []}

    # Garble replaces function names with hashes
    hash_names = sum(1 for s in strings_list if re.match(r"^[a-f0-9]{16,}$", s))
    if hash_names > 20:
        indicators["obfuscated"] = True
        indicators["techniques"].append("Possible garble obfuscation (hash-like function names)")

    # Missing gopclntab suggests stripping
    if not go_info.get("gopclntab_found"):
        indicators["techniques"].append("gopclntab not found - may be stripped or modified")

    # Low runtime string count
    if go_info.get("runtime_strings_found", 0) < 2:
        indicators["obfuscated"] = True
        indicators["techniques"].append("Low Go runtime string count - possible obfuscation")

    return indicators


def generate_report(filepath):
    """Generate comprehensive Go malware analysis report."""
    report = {
        "file": filepath,
        "sha256": compute_hash(filepath),
        "size": os.path.getsize(filepath),
    }

    go_info = detect_go_binary(filepath)
    report["go_detection"] = go_info

    if not go_info["is_go_binary"]:
        report["conclusion"] = "Not identified as a Go binary"
        return report

    strings_list = extract_go_strings(filepath)
    report["total_strings"] = len(strings_list)

    packages = extract_go_packages(strings_list)
    report["packages"] = packages[:50]

    suspicious = detect_suspicious_packages(packages)
    report["suspicious_packages"] = suspicious

    sections = analyze_sections(filepath)
    report["sections"] = sections

    obfuscation = detect_obfuscation(go_info, strings_list)
    report["obfuscation"] = obfuscation

    return report


if __name__ == "__main__":
    print("=" * 60)
    print("Go Malware Analysis Agent (Ghidra-assisted)")
    print("Go binary detection, package extraction, obfuscation detection")
    print("=" * 60)

    target = sys.argv[1] if len(sys.argv) > 1 else None

    if not target or not os.path.exists(target):
        print("\n[DEMO] Usage: python agent.py <go_binary>")
        sys.exit(0)

    report = generate_report(target)
    go = report.get("go_detection", {})
    print(f"\n[*] File: {target}")
    print(f"[*] SHA-256: {report['sha256']}")
    print(f"[*] Go binary: {go.get('is_go_binary', False)}")
    print(f"[*] Go version: {go.get('go_version', 'unknown')}")
    print(f"[*] Strings: {report.get('total_strings', 0)}")

    print("\n--- Packages ---")
    for pkg in report.get("packages", [])[:15]:
        print(f"  {pkg}")

    print("\n--- Suspicious Packages ---")
    for s in report.get("suspicious_packages", []):
        print(f"  [!] {s['package']}: {s['concern']}")

    print("\n--- Obfuscation ---")
    obf = report.get("obfuscation", {})
    print(f"  Obfuscated: {obf.get('obfuscated', False)}")
    for t in obf.get("techniques", []):
        print(f"  {t}")

    print(f"\n{json.dumps(report, indent=2, default=str)}")