"""
getDependenciesTreeImgDLL.py

Recursively analyzes DLL dependencies of a Windows PE file. If a DLL 
can dynamically load another DLL that DLL will not necessarily appear 
in the normal import table. This program is best treated as a static 
dependency analyzer, not as an exact reproduction of the Windows loader.

Outputs:
    <output>.dot
    <output>.png
    <output>_nodes.csv         - Contains 1 row for every discovered DLL
    <output>_dependencies.csv  - Provides actual dependency relationships
    <output>_missing.csv       - Contains only unresolved DLLs

Usages:
----------
For a DLL: 
    python dll_dependency_analyzer.py LibThirdParty.dll

Specify additional directories containing DLLs of application:
    python dll_dependency_analyzer.py LibThirdParty.dll -s ./lib -s ./etc

Specify the output prefix:
    python getDependenciesTreeImgDLL.py LibThirdParty.dll -o dep_graph
    
Features:
    - Missing DLL detection
    - x86/x64/ARM/ARM64 detection
    - Architecture mismatch detection
    - Circular dependency detection
    - Graphviz DOT generation
    - PNG generation, CSV export
    - Search through user supplied directories and PATH

The generated PNG marks:
Blue    → x64
Green   → x86
Purple  → ARM64
Yellow  → ARM
Red     → missing DLL
Orange  → architecture mismatch
"""

from __future__ import annotations

import argparse
import csv
import os
import shutil
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import pefile


# ---------------------------------------------------------------------------
# PE architecture definitions
# ---------------------------------------------------------------------------

MACHINE_TYPES = {
    0x014C: "x86",
    0x8664: "x64",
    0x01C0: "ARM",
    0xAA64: "ARM64",
    0x0200: "IA64",
}


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------

@dataclass
class DLLNode:
    name: str
    path: Optional[Path]
    architecture: str
    status: str
    depth: int


@dataclass
class DependencyEdge:
    parent: str
    child: str
    child_path: Optional[Path]
    child_architecture: str
    status: str


# ---------------------------------------------------------------------------
# PE utilities
# ---------------------------------------------------------------------------

def get_architecture(pe: pefile.PE) -> str:
    """
    Return the CPU architecture of a PE file.
    """

    machine = pe.FILE_HEADER.Machine

    return MACHINE_TYPES.get(
        machine,
        f"Unknown (0x{machine:04X})"
    )


def get_pe_architecture(path: Path) -> str:
    """
    Read PE architecture without executing the file.
    """

    try:
        pe = pefile.PE(str(path), fast_load=True)
        architecture = get_architecture(pe)
        pe.close()

        return architecture

    except Exception:
        return "Invalid PE"


def get_imports(path: Path) -> list[str]:
    """
    Return DLL names from the PE import table.
    """

    imports = []

    try:
        pe = pefile.PE(str(path), fast_load=True)

        pe.parse_data_directories(
            directories=[
                pefile.DIRECTORY_ENTRY[
                    "IMAGE_DIRECTORY_ENTRY_IMPORT"
                ]
            ]
        )

        if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):

            for entry in pe.DIRECTORY_ENTRY_IMPORT:

                if entry.dll:

                    name = entry.dll.decode(
                        "utf-8",
                        errors="replace"
                    )

                    imports.append(name)

        pe.close()

    except Exception as exc:

        print(
            f"WARNING: Could not read imports from "
            f"{path}: {exc}"
        )

    return sorted(set(imports))


# ---------------------------------------------------------------------------
# Windows search paths
# ---------------------------------------------------------------------------

def get_default_search_paths() -> list[Path]:
    """
    Build a list of directories that are commonly searched
    for Windows DLLs.
    """

    directories = []

    # Current directory
    directories.append(Path.cwd())

    # Windows directory
    windir = os.environ.get("WINDIR")

    if windir:

        windows_dir = Path(windir)

        directories.append(windows_dir)
        directories.append(windows_dir / "System32")
        directories.append(windows_dir / "SysWOW64")

    # PATH
    path_variable = os.environ.get("PATH", "")

    for directory in path_variable.split(os.pathsep):

        if directory:

            directories.append(Path(directory))

    # Remove duplicates
    unique = []

    seen = set()

    for directory in directories:

        try:
            directory = directory.resolve()
        except Exception:
            continue

        key = str(directory).lower()

        if key not in seen:

            seen.add(key)
            unique.append(directory)

    return unique


# ---------------------------------------------------------------------------
# DLL locator
# ---------------------------------------------------------------------------

class DLLLocator:

    def __init__(self, search_paths: list[Path]):

        self.search_paths = search_paths

        # Cache avoids repeatedly searching the same DLL
        self.cache: dict[str, Optional[Path]] = {}

    def find(self, dll_name: str) -> Optional[Path]:

        key = dll_name.lower()

        if key in self.cache:
            return self.cache[key]

        # Remove possible surrounding quotes
        dll_name = dll_name.strip('"')

        # Absolute path
        candidate = Path(dll_name)

        if candidate.is_absolute():

            if candidate.exists():

                self.cache[key] = candidate.resolve()

                return self.cache[key]

        # Search directories
        for directory in self.search_paths:

            candidate = directory / dll_name

            if candidate.exists() and candidate.is_file():

                result = candidate.resolve()

                self.cache[key] = result

                return result

        self.cache[key] = None

        return None


# ---------------------------------------------------------------------------
# Dependency analyzer
# ---------------------------------------------------------------------------

class DependencyAnalyzer:

    def __init__(
        self,
        root: Path,
        search_paths: list[Path],
    ):

        self.root = root.resolve()

        self.root_architecture = get_pe_architecture(
            self.root
        )

        self.locator = DLLLocator(search_paths)

        self.nodes: dict[str, DLLNode] = {}

        self.edges: list[DependencyEdge] = []

        self.missing: set[str] = set()

        self.cycles: list[tuple[str, str]] = []

        # Keep track of recursion stack
        self.recursion_stack: list[str] = []

        # Keep track of analyzed files
        self.analyzed: set[str] = set()

    # ---------------------------------------------------------------------

    def node_key(self, path: Path) -> str:

        return str(path.resolve()).lower()

    # ---------------------------------------------------------------------

    def architecture_mismatch(
        self,
        parent_arch: str,
        child_arch: str,
    ) -> bool:

        # Unknown architectures cannot reliably be compared
        if parent_arch.startswith("Unknown"):
            return False

        if child_arch.startswith("Unknown"):
            return False

        if child_arch == "Invalid PE":
            return False

        # ARM and x86/x64 etc. are incompatible
        return parent_arch != child_arch

    # ---------------------------------------------------------------------

    def analyze(self):

        print()
        print("=" * 70)
        print("DLL DEPENDENCY ANALYSIS")
        print("=" * 70)

        print(f"Root DLL : {self.root}")
        print(f"Architecture : {self.root_architecture}")

        self._visit(
            self.root,
            depth=0,
            parent=None,
            parent_arch=None,
        )

    # ---------------------------------------------------------------------

    def _visit(
        self,
        path: Path,
        depth: int,
        parent: Optional[Path],
        parent_arch: Optional[str],
    ):

        path = path.resolve()

        key = self.node_key(path)

        architecture = get_pe_architecture(path)

        # Register node
        if key not in self.nodes:

            self.nodes[key] = DLLNode(
                name=path.name,
                path=path,
                architecture=architecture,
                status="found",
                depth=depth,
            )

        # Circular dependency
        if key in self.recursion_stack:

            if parent:

                self.cycles.append(
                    (
                        parent.name,
                        path.name
                    )
                )

            return

        # Already analyzed
        if key in self.analyzed:
            return

        self.recursion_stack.append(key)

        imports = get_imports(path)

        print(
            f"{'  ' * depth}"
            f"{path.name} "
            f"[{architecture}]"
        )

        for dependency in imports:

            dependency_path = self.locator.find(
                dependency
            )

            if dependency_path is None:

                print(
                    f"{'  ' * (depth + 1)}"
                    f"[MISSING] {dependency}"
                )

                self.missing.add(dependency.lower())

                edge = DependencyEdge(
                    parent=path.name,
                    child=dependency,
                    child_path=None,
                    child_architecture="MISSING",
                    status="missing",
                )

                self.edges.append(edge)

                continue

            dependency_arch = get_pe_architecture(
                dependency_path
            )

            mismatch = self.architecture_mismatch(
                architecture,
                dependency_arch,
            )

            if mismatch:

                status = "ARCHITECTURE MISMATCH"

                print(
                    f"{'  ' * (depth + 1)}"
                    f"[{status}] "
                    f"{dependency} "
                    f"({dependency_arch})"
                )

            else:

                status = "ok"

                print(
                    f"{'  ' * (depth + 1)}"
                    f"{dependency} "
                    f"({dependency_arch})"
                )

            edge = DependencyEdge(
                parent=path.name,
                child=dependency_path.name,
                child_path=dependency_path,
                child_architecture=dependency_arch,
                status=status,
            )

            self.edges.append(edge)

            # Register dependency node
            dependency_key = self.node_key(
                dependency_path
            )

            if dependency_key not in self.nodes:

                self.nodes[dependency_key] = DLLNode(
                    name=dependency_path.name,
                    path=dependency_path,
                    architecture=dependency_arch,
                    status=status,
                    depth=depth + 1,
                )

            # Recursive analysis
            self._visit(
                dependency_path,
                depth + 1,
                path,
                architecture,
            )

        self.recursion_stack.pop()

        self.analyzed.add(key)


# ---------------------------------------------------------------------------
# DOT generation
# ---------------------------------------------------------------------------

def escape_dot(text: str) -> str:

    return (
        text
        .replace("\\", "\\\\")
        .replace('"', '\\"')
    )


def node_color(status: str) -> str:

    if status == "missing":
        return "red"

    if status == "ARCHITECTURE MISMATCH":
        return "orange"

    return "lightblue"


def architecture_color(architecture: str) -> str:

    if architecture == "x64":
        return "lightblue"

    if architecture == "x86":
        return "lightgreen"

    if architecture == "ARM64":
        return "plum"

    if architecture == "ARM":
        return "khaki"

    return "gray"


def generate_dot(
    analyzer: DependencyAnalyzer,
    output: Path,
):

    lines = []

    lines.append("digraph DLLDependencies {")

    lines.append(
        '    rankdir=LR;'
    )

    lines.append(
        '    graph [fontname="Arial"];'
    )

    lines.append(
        '    node [shape=box, style="rounded,filled", '
        'fontname="Arial"];'
    )

    lines.append(
        '    edge [fontname="Arial"];'
    )

    # ------------------------------------------------------------------
    # Nodes
    # ------------------------------------------------------------------

    for key, node in analyzer.nodes.items():

        label = (
            f"{node.name}\\n"
            f"Architecture: {node.architecture}"
        )

        if node.status == "ARCHITECTURE MISMATCH":
            fill = "orange"
        else:
            fill = architecture_color(
                node.architecture
            )

        node_id = "n" + str(
            abs(hash(key))
        )

        lines.append(
            f'    {node_id} '
            f'[label="{escape_dot(label)}", '
            f'fillcolor="{fill}"];'
        )

    # Missing nodes
    missing_ids = {}

    for missing in sorted(analyzer.missing):

        node_id = "missing_" + str(
            abs(hash(missing))
        )

        missing_ids[missing] = node_id

        lines.append(
            f'    {node_id} '
            f'[label="{escape_dot(missing)}\\nMISSING", '
            f'fillcolor="red", '
            f'fontcolor="white"];'
        )

    # ------------------------------------------------------------------
    # Edges
    # ------------------------------------------------------------------

    path_to_id = {}

    for key in analyzer.nodes:

        path_to_id[key] = "n" + str(
            abs(hash(key))
        )

    for edge in analyzer.edges:

        parent_path = None

        for key, node in analyzer.nodes.items():

            if node.name.lower() == edge.parent.lower():

                parent_path = key
                break

        if parent_path is None:
            continue

        parent_id = path_to_id[parent_path]

        if edge.status == "missing":

            child_id = missing_ids[
                edge.child.lower()
            ]

            edge_color = "red"
            style = "dashed"

        else:

            child_path = (
                edge.child_path.resolve()
                if edge.child_path
                else None
            )

            if child_path:

                child_key = str(
                    child_path
                ).lower()

                child_id = path_to_id.get(
                    child_key
                )

                if child_id is None:
                    continue

            else:
                continue

            if edge.status == "ARCHITECTURE MISMATCH":

                edge_color = "orange"
                style = "bold"

            else:

                edge_color = "black"
                style = "solid"

        lines.append(
            f'    {parent_id} -> {child_id} '
            f'[color="{edge_color}", '
            f'style="{style}"];'
        )

    lines.append("}")

    output.write_text(
        "\n".join(lines),
        encoding="utf-8",
    )


# ---------------------------------------------------------------------------
# Graphviz PNG
# ---------------------------------------------------------------------------

def generate_png(
    dot_file: Path,
    png_file: Path,
):

    dot_executable = shutil.which("dot")

    if dot_executable is None:

        print()
        print(
            "WARNING: Graphviz 'dot' executable was not found."
        )

        print(
            "DOT file was generated, but PNG was not."
        )

        return False

    try:

        subprocess.run(
            [
                dot_executable,
                "-Tpng",
                str(dot_file),
                "-o",
                str(png_file),
            ],
            check=True,
        )

        return True

    except subprocess.CalledProcessError as exc:

        print(
            f"ERROR: Graphviz failed: {exc}"
        )

        return False


# ---------------------------------------------------------------------------
# CSV export
# ---------------------------------------------------------------------------

def export_nodes_csv(
    analyzer: DependencyAnalyzer,
    output: Path,
):

    with output.open(
        "w",
        newline="",
        encoding="utf-8",
    ) as file:

        writer = csv.writer(file)

        writer.writerow(
            [
                "DLL",
                "Path",
                "Architecture",
                "Status",
                "Depth",
            ]
        )

        for node in analyzer.nodes.values():

            writer.writerow(
                [
                    node.name,
                    str(node.path)
                    if node.path
                    else "",
                    node.architecture,
                    node.status,
                    node.depth,
                ]
            )


def export_dependencies_csv(
    analyzer: DependencyAnalyzer,
    output: Path,
):

    with output.open(
        "w",
        newline="",
        encoding="utf-8",
    ) as file:

        writer = csv.writer(file)

        writer.writerow(
            [
                "Parent DLL",
                "Dependency",
                "Dependency Path",
                "Dependency Architecture",
                "Status",
            ]
        )

        for edge in analyzer.edges:

            writer.writerow(
                [
                    edge.parent,
                    edge.child,
                    str(edge.child_path)
                    if edge.child_path
                    else "",
                    edge.child_architecture,
                    edge.status,
                ]
            )


def export_missing_csv(
    analyzer: DependencyAnalyzer,
    output: Path,
):

    with output.open(
        "w",
        newline="",
        encoding="utf-8",
    ) as file:

        writer = csv.writer(file)

        writer.writerow(
            [
                "Missing DLL"
            ]
        )

        for dll in sorted(analyzer.missing):

            writer.writerow(
                [
                    dll
                ]
            )


# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------

def print_summary(
    analyzer: DependencyAnalyzer,
):

    mismatch_count = sum(
        1
        for edge in analyzer.edges
        if edge.status == "ARCHITECTURE MISMATCH"
    )

    print()
    print("=" * 70)
    print("SUMMARY")
    print("=" * 70)

    print(
        f"Root architecture : "
        f"{analyzer.root_architecture}"
    )

    print(
        f"DLLs found        : "
        f"{len(analyzer.nodes)}"
    )

    print(
        f"Dependencies      : "
        f"{len(analyzer.edges)}"
    )

    print(
        f"Missing DLLs      : "
        f"{len(analyzer.missing)}"
    )

    print(
        f"Architecture      : "
        f"{mismatch_count}"
    )

    print(
        f"Circular references: "
        f"{len(analyzer.cycles)}"
    )

    if analyzer.missing:

        print()
        print("Missing DLLs:")

        for dll in sorted(analyzer.missing):

            print(
                f"  - {dll}"
            )

    if analyzer.cycles:

        print()
        print("Circular dependencies:")

        for parent, child in analyzer.cycles:

            print(
                f"  {parent} -> {child}"
            )


# ---------------------------------------------------------------------------
# Command line
# ---------------------------------------------------------------------------

def main():

    parser = argparse.ArgumentParser(
        description=(
            "Recursively analyze DLL dependencies "
            "and generate Graphviz/CSV output."
        )
    )

    parser.add_argument(
        "dll",
        help="Root DLL or EXE to analyze",
    )

    parser.add_argument(
        "-o",
        "--output",
        default="dll_dependency",
        help=(
            "Output file prefix "
            "(default: dll_dependency)"
        ),
    )

    parser.add_argument(
        "-s",
        "--search",
        action="append",
        default=[],
        help=(
            "Additional DLL search directory. "
            "Can be specified multiple times."
        ),
    )

    parser.add_argument(
        "--no-system-paths",
        action="store_true",
        help=(
            "Do not search Windows directories/PATH."
        ),
    )

    args = parser.parse_args()

    root = Path(args.dll).resolve()

    if not root.exists():

        print(
            f"ERROR: File does not exist: {root}"
        )

        sys.exit(1)

    # ------------------------------------------------------------------
    # Build search paths
    # ------------------------------------------------------------------

    search_paths = []

    # Always search directory containing root DLL
    search_paths.append(root.parent)

    if not args.no_system_paths:

        search_paths.extend(
            get_default_search_paths()
        )

    # User supplied directories
    for directory in args.search:

        search_paths.append(
            Path(directory).resolve()
        )

    # Remove duplicates
    unique_paths = []

    seen = set()

    for directory in search_paths:

        key = str(directory).lower()

        if key not in seen:

            seen.add(key)
            unique_paths.append(directory)

    # ------------------------------------------------------------------
    # Analyze
    # ------------------------------------------------------------------

    analyzer = DependencyAnalyzer(
        root=root,
        search_paths=unique_paths,
    )

    analyzer.analyze()

    # ------------------------------------------------------------------
    # Output names
    # ------------------------------------------------------------------

    output_prefix = Path(args.output)

    dot_file = output_prefix.with_suffix(".dot")

    png_file = output_prefix.with_suffix(".png")

    nodes_csv = Path(
        str(output_prefix) + "_nodes.csv"
    )

    dependencies_csv = Path(
        str(output_prefix) + "_dependencies.csv"
    )

    missing_csv = Path(
        str(output_prefix) + "_missing.csv"
    )

    # ------------------------------------------------------------------
    # Generate outputs
    # ------------------------------------------------------------------

    generate_dot(
        analyzer,
        dot_file,
    )

    generate_png(
        dot_file,
        png_file,
    )

    export_nodes_csv(
        analyzer,
        nodes_csv,
    )

    export_dependencies_csv(
        analyzer,
        dependencies_csv,
    )

    export_missing_csv(
        analyzer,
        missing_csv,
    )

    # ------------------------------------------------------------------
    # Summary
    # ------------------------------------------------------------------

    print_summary(analyzer)

    print()
    print("Output files:")
    print(f"  DOT          : {dot_file}")
    print(f"  PNG          : {png_file}")
    print(f"  Nodes CSV    : {nodes_csv}")
    print(f"  Dependencies : {dependencies_csv}")
    print(f"  Missing CSV  : {missing_csv}")


if __name__ == "__main__":
    main()