import logging
from rich.logging import RichHandler
from rich.console import Console
from rich.progress import track

import argparse
import os, pathlib
from pickle import Pickler, Unpickler

from bloodhound.ad.utils import ADUtils

from collections import defaultdict
from requests.structures import CaseInsensitiveDict

import datetime
from enum import Enum

class ADExplorerSnapshot(object):
    OutputMode = Enum('OutputMode', ['BOFHound', 'BloodHound', 'Objects'])

    def __init__(self, snapfile, outputfolder, console=None, snapshot_parser=None):
        self.console = console or setup_logging()
        self.output = outputfolder

        if not snapshot_parser:
            from adexpsnapshot.parser.classes import Snapshot
            snapshot_parser = Snapshot

        self.snapfile = snapfile
        self.snap = snapshot_parser(snapfile)

        self.snap.parseHeader()

        filetimeiso = datetime.datetime.fromtimestamp(self.snap.header.filetimeUnix).isoformat()
        logging.info(f'Server: {self.snap.header.server}')
        logging.info(f'Time of snapshot: {filetimeiso}')
        logging.info('Metadata offset: 0x{:x}'.format(self.snap.header.metadataOffset))
        logging.info(f'Object count: {self.snap.header.numObjects}')

        self.snap.parseProperties()
        self.snap.parseClasses()
        self.snap.parseObjectOffsets()

        self.sidcache = {}
        self.dncache = CaseInsensitiveDict()
        self.computersidcache = CaseInsensitiveDict()
        self.domains = CaseInsensitiveDict()
        self.objecttype_guid_map = CaseInsensitiveDict()
        self.domaincontrollers = []
        self.certtemplates = defaultdict(set)

        self.rootdomain = None
        self.domain_dn = None
        self.config_dn = None
        self.schema_dn = None
        self.forest_dn = None

    def preprocess(self, cache=False):
        if cache:
            cacheFileName = str(pathlib.Path(self.snapfile.name).with_suffix('.cache'))

            dico = None
            try:
                dico = Unpickler(open(cacheFileName, "rb")).load()
            except (OSError, IOError, EOFError) as e:
                pass

            if dico and dico.get('unixtime') == self.snap.header.filetimeUnix:
                self.objecttype_guid_map = dico['guidmap']
                self.sidcache = dico['sidcache']
                self.dncache = dico['dncache']
                self.computersidcache = dico['computersidcache']
                self.domains = dico['domains']
                self.domaincontrollers = dico['domaincontrollers']
                self.certtemplates = dico['certtemplates']

                self.rootdomain = dico['rootdomain']
                self.domain_dn = dico['domain_dn']
                self.config_dn = dico['config_dn']
                self.schema_dn = dico['schema_dn']
                self.forest_dn = dico['forest_dn']

                self.console.print(f"[green]✓[/green] Using cached data for preprocessing")
                return

        for k, cl in self.snap.classes.items():
            self.objecttype_guid_map[k] = str(cl.schemaIDGUID)

        for k, idx in self.snap.propertyDict.items():
            self.objecttype_guid_map[k] = str(self.snap.properties[idx].schemaIDGUID)

        for idx, obj in track(enumerate(self.snap.objects), description="Preprocessing", total=self.snap.header.numObjects):
            # create sid cache
            objectSid = ADUtils.get_entry_property(obj, 'objectSid')
            if objectSid:
                self.sidcache[str(objectSid)] = idx

            # create dn cache
            distinguishedName = ADUtils.get_entry_property(obj, 'distinguishedName')
            if distinguishedName:
                self.dncache[str(distinguishedName)] = idx

            # get domains
            # Domain objects have 'domain' class AND objectSid attribute
            # DNS zones have 'domain' class but no objectSid
            if 'domain' in obj.classes and objectSid:
                if self.rootdomain is not None:  # is it possible to find multiple?
                    logging.warning("Multiple domains in snapshot(?)")
                else:
                    self.rootdomain = str(distinguishedName) # keeping for backward compat
                    self.domain_dn = str(distinguishedName)
                    self.domains[str(distinguishedName)] = idx

                    objectCategory = ADUtils.get_entry_property(obj, 'objectCategory')
                    try:
                        _, self.schema_dn = objectCategory.split(',', 1)
                        _, self.config_dn = self.schema_dn.split(',', 1)
                        _, self.forest_dn = self.config_dn.split(',', 1)
                    except (ValueError, AttributeError):
                        logging.warning("Failed to parse schema, config, or forest DN from objectCategory")
                        pass

            # get forest domains
            if 'crossref' in obj.classes:
                if ADUtils.get_entry_property(obj, 'systemFlags', 0) & 2 == 2:
                    ncname = ADUtils.get_entry_property(obj, 'nCName')
                    if ncname and ncname not in self.domains:
                        self.domains[str(ncname)] = idx

            # get computers
            if ADUtils.get_entry_property(obj, 'sAMAccountType', -1) == 805306369:
                dnshostname = ADUtils.get_entry_property(obj, 'dNSHostname')
                if dnshostname:
                    self.computersidcache[str(dnshostname)] = str(objectSid)

            # get all cert templates
            if 'pkienrollmentservice' in obj.classes:
                name = str(ADUtils.get_entry_property(obj, 'name'))
                if ADUtils.get_entry_property(obj, 'certificateTemplates'):
                    templates = ADUtils.get_entry_property(obj, 'certificateTemplates')
                    for template in templates:
                        self.certtemplates[str(template)].add(name)

            # get dcs
            if ADUtils.get_entry_property(obj, 'userAccountControl', 0) & 0x2000 == 0x2000:
                self.domaincontrollers.append(idx)

        if cache:
            dico = {
                'guidmap': self.objecttype_guid_map,
                'sidcache': self.sidcache,
                'dncache': self.dncache,
                'computersidcache': self.computersidcache,
                'domains': self.domains,
                'domaincontrollers': self.domaincontrollers,
                'rootdomain': self.rootdomain,
                'domain_dn': self.domain_dn,
                'config_dn': self.config_dn,
                'schema_dn': self.schema_dn,
                'forest_dn': self.forest_dn,
                'certtemplates': self.certtemplates,
                'unixtime': self.snap.header.filetimeUnix
            }
            with open(cacheFileName, "wb") as f:
                Pickler(f).dump(dico)
            self.console.print(f"[green]✓[/green] Cache saved to {cacheFileName}")
        
        self.console.print(f"[green]✓[/green] Preprocessing complete")

    def preprocessCached(self):
        """Alias for preprocess() for backward compatibility with scripts."""
        self.preprocess(cache=True)

    def outputObjects(self):
        """Output all objects to NDJSON file."""
        from adexpsnapshot.ouput.objects import ObjectsOutput
        handler = ObjectsOutput(self.snap, self.output, self.console)
        handler.process()

    def outputBOFHound(self):
        """Output objects in BOFHound log format."""
        from adexpsnapshot.ouput.bofhound import BOFHoundOutput
        handler = BOFHoundOutput(self.snap, self.output, self.console)
        handler.process()

    def outputBloodHound(self):
        """Output data in BloodHound JSON format."""

        logging.warning("The BloodHound mode gives incomplete output. Please use the BOFHound mode instead.")

        self.preprocess(cache=True)
        from adexpsnapshot.ouput.bloodhound import BloodHoundOutput
        handler = BloodHoundOutput(
            self.snap, self.output, self.console,
            self.sidcache, self.dncache, self.computersidcache, self.domains,
            self.objecttype_guid_map, self.domaincontrollers, self.rootdomain, self.certtemplates
        )
        handler.process()


def setup_logging(level=logging.INFO):
    console = Console()
    rich_handler = RichHandler(
        console=console,
        show_time=True,
        show_path=False,
        rich_tracebacks=True,
        markup=True
    )
    
    logging.basicConfig(
        level=level,
        format="%(message)s",
        datefmt="[%X]",
        handlers=[rich_handler]
    )
    
    return console

def mode_output(ades, args):
    """Handle output mode operations (BOFHound, BloodHound, Objects)."""
    outputmode = ADExplorerSnapshot.OutputMode[args.mode]
    if outputmode == ADExplorerSnapshot.OutputMode.BloodHound:
        ades.outputBloodHound()
    elif outputmode == ADExplorerSnapshot.OutputMode.Objects:
        ades.outputObjects()
    elif outputmode == ADExplorerSnapshot.OutputMode.BOFHound:
        ades.outputBOFHound()

def mode_enrich(ades):
    """Handle enrichment mode - reconstruct missing treeview metadata."""
    from adexpsnapshot.enrich import enrich_snapshot
    res = enrich_snapshot(ades)
    if res:
        ades.console.print(f"[green]✓[/green] Enrichment complete")
    else:
        ades.console.print(f"[red]✗[/red] Enrichment skipped or failed")

def main():
    console = setup_logging()

    parser = argparse.ArgumentParser(add_help=True, description='AD Explorer snapshot ingestor for BloodHound', formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument('snapshot', type=argparse.FileType('rb'), help="Path to the snapshot .dat file.")
    parser.add_argument('-o', '--output', required=False, type=pathlib.Path, help="Path to an output folder. Folder will be created if it doesn't exist. Defaults to the current directory.", default=".")
    parser.add_argument('-m', '--mode', required=False, help="The output mode to use. Defaults to BOFHound output mode, which can then be used with BOFHound. Can also directly output to BloodHound JSON output files (with limitations). In Objects mode all objects with all attributes are outputted to NDJSON.", choices=ADExplorerSnapshot.OutputMode.__members__, default='BOFHound')
    parser.add_argument('-e', '--enrich', action='store_true', help="Reconstruct treeview metadata if missing and save to enriched file.")

    args = parser.parse_args()

    if not os.path.exists(args.output):
        try:
            os.mkdir(args.output)
        except:
            logging.error(f"Unable to create output directory '{args.output}'.")
            return
    
    if not os.path.isdir(args.output):
        logging.error(f"Path '{args.output}' does not exist or is not a folder.")
        return
    
    ades = ADExplorerSnapshot(args.snapshot, args.output, console=console)
    
    if args.enrich:
        mode_enrich(ades)
    else:
        mode_output(ades, args)

if __name__ == '__main__':
    main()
