"""Build a comparable file inventory for JxVoCongTruyenKy baseline."""
from __future__ import annotations

import csv
import hashlib
import json
import os
import sqlite3
import datetime
from pathlib import Path

ROOT = Path(r"C:\JXVCTK\JxVoCongTruyenKy")
OUT = Path(r"C:\Projects\JX\security-baseline\2026-09-18-vctk")

PE_EXTS = {".exe", ".dll", ".sys", ".asi", ".drv"}


def sha256_md5(path: Path, chunk: int = 1024 * 1024) -> tuple[str, str]:
    h256 = hashlib.sha256()
    hmd5 = hashlib.md5()
    with path.open("rb") as f:
        while True:
            b = f.read(chunk)
            if not b:
                break
            h256.update(b)
            hmd5.update(b)
    return h256.hexdigest(), hmd5.hexdigest()


def pe_info(path: Path) -> dict:
    info = {
        "is_pe": False,
        "machine": None,
        "ep_rva": None,
        "image_base": None,
        "sections": None,
        "imports": None,
        "entropy_hint": None,
    }
    try:
        import pefile
        import collections
        import math

        pe = pefile.PE(str(path), fast_load=True)
        info["is_pe"] = True
        info["machine"] = hex(pe.FILE_HEADER.Machine)
        info["ep_rva"] = hex(pe.OPTIONAL_HEADER.AddressOfEntryPoint)
        info["image_base"] = hex(pe.OPTIONAL_HEADER.ImageBase)
        secs = []
        ents = []
        for s in pe.sections:
            name = s.Name.split(b"\x00", 1)[0].decode("latin1", "replace")
            data = s.get_data()[:65536]
            if data:
                c = collections.Counter(data)
                n = len(data)
                ent = -sum((v / n) * math.log2(v / n) for v in c.values())
            else:
                ent = 0.0
            ents.append(ent)
            secs.append(f"{name}:{ent:.2f}")
        info["sections"] = ";".join(secs)
        info["entropy_hint"] = round(max(ents) if ents else 0.0, 2)
        try:
            pe.parse_data_directories(
                directories=[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IMPORT"]]
            )
            imps = []
            if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
                for e in pe.DIRECTORY_ENTRY_IMPORT:
                    dll = e.dll.decode("latin1", "replace") if e.dll else "?"
                    names = []
                    for imp in e.imports[:40]:
                        names.append(
                            imp.name.decode("latin1", "replace")
                            if imp.name
                            else f"ord_{imp.ordinal}"
                        )
                    imps.append(dll + "=" + ",".join(names))
            info["imports"] = " | ".join(imps)[:4000]
        except Exception:
            pass
        pe.close()
    except Exception as ex:
        info["pe_error"] = str(ex)
    return info


def main() -> None:
    rows = []
    for dirpath, dirnames, filenames in os.walk(ROOT):
        dirnames[:] = [d for d in dirnames if d not in {".git"}]
        for name in filenames:
            p = Path(dirpath) / name
            rel = str(p.relative_to(ROOT)).replace("\\", "/")
            st = p.stat()
            rec = {
                "relpath": rel,
                "size": st.st_size,
                "mtime_utc": datetime.datetime.utcfromtimestamp(st.st_mtime).isoformat()
                + "Z",
                "ctime_utc": datetime.datetime.utcfromtimestamp(st.st_ctime).isoformat()
                + "Z",
                "ext": p.suffix.lower(),
                "sha256": None,
                "md5": None,
            }
            try:
                rec["sha256"], rec["md5"] = sha256_md5(p)
            except OSError as ex:
                rec["hash_error"] = str(ex)
                rec["sha256"] = ""
                rec["md5"] = ""
            if p.suffix.lower() in PE_EXTS and rec.get("sha256"):
                rec.update(pe_info(p))
            else:
                rec.update(
                    {
                        "is_pe": False,
                        "machine": None,
                        "ep_rva": None,
                        "image_base": None,
                        "sections": None,
                        "imports": None,
                        "entropy_hint": None,
                    }
                )
            rows.append(rec)
            print(f"{rec['sha256'][:12] if rec['sha256'] else 'ERR':12} {rec['size']:10} {rel}")

    rows.sort(key=lambda r: r["relpath"].lower())

    # CSV
    csv_path = OUT / "inventory.csv"
    fields = [
        "relpath",
        "size",
        "sha256",
        "md5",
        "ext",
        "mtime_utc",
        "ctime_utc",
        "is_pe",
        "machine",
        "ep_rva",
        "image_base",
        "entropy_hint",
        "sections",
        "imports",
    ]
    with csv_path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore")
        w.writeheader()
        w.writerows(rows)

    # SHA256SUMS
    with (OUT / "SHA256SUMS.txt").open("w", encoding="utf-8") as f:
        for r in rows:
            if r.get("sha256"):
                f.write(f"{r['sha256']}  {r['relpath']}\n")

    # JSON
    manifest = {
        "snapshot_id": "2026-09-18-vctk",
        "created_local": "2026-09-18T20:30:00+07:00",
        "source_root": str(ROOT),
        "file_count": len(rows),
        "total_bytes": sum(r["size"] for r in rows),
        "files": rows,
    }
    with (OUT / "inventory.json").open("w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)

    # SQLite
    db = OUT / "inventory.sqlite"
    if db.exists():
        db.unlink()
    con = sqlite3.connect(db)
    cur = con.cursor()
    cur.execute(
        """
        CREATE TABLE files (
            relpath TEXT PRIMARY KEY,
            size INTEGER,
            sha256 TEXT,
            md5 TEXT,
            ext TEXT,
            mtime_utc TEXT,
            ctime_utc TEXT,
            is_pe INTEGER,
            machine TEXT,
            ep_rva TEXT,
            image_base TEXT,
            entropy_hint REAL,
            sections TEXT,
            imports TEXT
        )
        """
    )
    cur.execute("CREATE INDEX idx_sha256 ON files(sha256)")
    cur.execute("CREATE INDEX idx_md5 ON files(md5)")
    cur.execute("CREATE INDEX idx_ext ON files(ext)")
    for r in rows:
        cur.execute(
            """INSERT INTO files VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)""",
            (
                r["relpath"],
                r["size"],
                r.get("sha256"),
                r.get("md5"),
                r.get("ext"),
                r.get("mtime_utc"),
                r.get("ctime_utc"),
                1 if r.get("is_pe") else 0,
                r.get("machine"),
                r.get("ep_rva"),
                r.get("image_base"),
                r.get("entropy_hint"),
                r.get("sections"),
                r.get("imports"),
            ),
        )
    con.commit()
    con.close()
    print(f"Wrote {len(rows)} rows -> {db}")


if __name__ == "__main__":
    main()