#!/usr/bin/env python3
"""Agent for analyzing Windows LNK shortcut files for forensic artifacts."""

import os
import json
import csv
import argparse
from datetime import datetime

import LnkParse3


def parse_lnk_file(filepath):
    """Parse a single LNK file and extract forensic artifacts."""
    with open(filepath, "rb") as f:
        lnk = LnkParse3.lnk_file(f)
        info = lnk.get_json()

    parsed = {
        "lnk_file": os.path.basename(filepath),
        "target_path": "",
        "working_dir": "",
        "arguments": "",
        "target_created": "",
        "target_modified": "",
        "target_accessed": "",
        "file_size": "",
        "drive_type": "",
        "volume_serial": "",
        "volume_label": "",
        "machine_id": "",
        "mac_address": "",
    }

    header = info.get("header", {})
    parsed["target_created"] = str(header.get("creation_time", ""))
    parsed["target_modified"] = str(header.get("modified_time", ""))
    parsed["target_accessed"] = str(header.get("accessed_time", ""))
    parsed["file_size"] = str(header.get("file_size", ""))

    link_info = info.get("link_info", {})
    if link_info:
        local_path = link_info.get("local_base_path", "")
        net_link = link_info.get("common_network_relative_link", {})
        network_path = net_link.get("net_name", "") if net_link else ""
        parsed["target_path"] = local_path or network_path

        vol_info = link_info.get("volume_id", {})
        if vol_info:
            parsed["drive_type"] = str(vol_info.get("drive_type", ""))
            parsed["volume_serial"] = str(vol_info.get("drive_serial_number", ""))
            parsed["volume_label"] = str(vol_info.get("volume_label", ""))

    string_data = info.get("string_data", {})
    parsed["working_dir"] = str(string_data.get("working_dir", ""))
    parsed["arguments"] = str(string_data.get("command_line_arguments", ""))

    extra = info.get("extra", {})
    tracker = extra.get("DISTRIBUTED_LINK_TRACKER_BLOCK", {})
    if tracker:
        parsed["machine_id"] = str(tracker.get("machine_id", ""))
        parsed["mac_address"] = str(tracker.get("mac_address", ""))

    return parsed


def parse_lnk_directory(directory):
    """Parse all LNK files in a directory."""
    results = []
    for filename in sorted(os.listdir(directory)):
        if not filename.lower().endswith(".lnk"):
            continue
        filepath = os.path.join(directory, filename)
        try:
            parsed = parse_lnk_file(filepath)
            results.append(parsed)
        except Exception as e:
            print(f"  Error parsing {filename}: {e}")
    return results


def filter_removable_media(results):
    """Filter LNK files that point to removable media."""
    return [r for r in results if "removable" in r.get("drive_type", "").lower()]


def filter_network_shares(results):
    """Filter LNK files pointing to network shares."""
    return [
        r for r in results
        if "network" in r.get("drive_type", "").lower()
        or r.get("target_path", "").startswith("\\\\")
    ]


def detect_suspicious_startup(startup_dir):
    """Analyze Startup folder LNK files for potential persistence."""
    suspicious = []
    for filename in os.listdir(startup_dir):
        if not filename.lower().endswith(".lnk"):
            continue
        filepath = os.path.join(startup_dir, filename)
        try:
            parsed = parse_lnk_file(filepath)
            target = parsed["target_path"].lower()
            args = parsed["arguments"].lower()
            if any(s in target for s in ["temp", "appdata", "programdata", "public"]):
                parsed["risk"] = "HIGH"
                suspicious.append(parsed)
            elif any(s in args for s in ["-enc", "powershell", "cmd /c", "wscript"]):
                parsed["risk"] = "HIGH"
                suspicious.append(parsed)
        except Exception:
            pass
    return suspicious


def export_csv(results, output_path):
    """Export parsed LNK results to CSV."""
    if not results:
        return
    with open(output_path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=results[0].keys())
        writer.writeheader()
        writer.writerows(results)


def extract_unique_machines(results):
    """Extract unique machine IDs and MAC addresses from LNK files."""
    machines = {}
    for r in results:
        mid = r.get("machine_id", "")
        mac = r.get("mac_address", "")
        if mid:
            machines[mid] = mac
    return machines


def main():
    parser = argparse.ArgumentParser(description="Windows LNK File Forensic Analysis Agent")
    parser.add_argument("--lnk-dir", required=True, help="Directory containing LNK files")
    parser.add_argument("--startup-dir", help="Startup folder to check for persistence")
    parser.add_argument("--output-dir", default="./lnk_analysis")
    parser.add_argument("--action", choices=[
        "parse_all", "removable", "network", "startup", "machines", "full_analysis"
    ], default="full_analysis")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    all_results = parse_lnk_directory(args.lnk_dir)
    print(f"[+] Parsed {len(all_results)} LNK files")

    if args.action in ("parse_all", "full_analysis"):
        csv_path = os.path.join(args.output_dir, "lnk_analysis.csv")
        export_csv(all_results, csv_path)
        print(f"[+] Exported to {csv_path}")

    if args.action in ("removable", "full_analysis"):
        removable = filter_removable_media(all_results)
        print(f"[+] Removable media files: {len(removable)}")
        for r in removable:
            print(f"    {r['target_modified']} | {r['target_path']} | Vol: {r['volume_serial']}")

    if args.action in ("network", "full_analysis"):
        network = filter_network_shares(all_results)
        print(f"[+] Network share files: {len(network)}")

    if args.action in ("startup", "full_analysis") and args.startup_dir:
        suspicious = detect_suspicious_startup(args.startup_dir)
        print(f"[+] Suspicious startup LNK: {len(suspicious)}")
        for s in suspicious:
            print(f"    [{s.get('risk')}] {s['lnk_file']} -> {s['target_path']}")

    if args.action in ("machines", "full_analysis"):
        machines = extract_unique_machines(all_results)
        print(f"[+] Unique machines: {len(machines)}")
        for mid, mac in machines.items():
            print(f"    Machine: {mid} | MAC: {mac}")

    print(json.dumps({"total_lnk": len(all_results), "generated_at": datetime.utcnow().isoformat()}, indent=2))


if __name__ == "__main__":
    main()