#!/usr/bin/env python3
"""
Sets up the victim precondition described in the report:

  "victim has at least one OIDC-managed Chrome profile (created previously
   via this enrollment flow) and is currently browsing in a DIFFERENT profile"

Concretely:
  1. Launch Chrome once with the Default profile to populate the user-data-dir.
  2. Launch Chrome once with --profile-directory="Profile 1" so Chrome creates
     and registers a second profile in Local State's profile.info_cache.
  3. Edit Local State to mark "Profile 1" as an OIDC-managed profile by setting
     the same attributes that the legitimate enrollment flow would have set
     (profile_management_oidc_auth_token, profile_management_oidc_id_token,
      profile_management_id = "iss:<iss>,sub:<sub>").

This is *precondition setup*, not part of the attack: it stands in for the
victim having previously gone through real Azure AD OIDC enrollment.
"""
import json
import os
import shutil
import subprocess
import sys
import time

CHROME = os.environ.get("CHROME", "/chromium/chromium/src/out/asan/chrome")
UDD = os.environ.get("UDD", "/tmp/chrome-poc-profile")

VICTIM_ISS = "https://login.microsoftonline.com/victim-tenant-id/v2.0"
VICTIM_SUB = "victim-pairwise-sub-abc123"
PROFILE_MGMT_ID = f"iss:{VICTIM_ISS},sub:{VICTIM_SUB}"


def run_chrome_once(profile_dir):
    """Run Chrome headless-ish so it creates/registers the profile, then exit."""
    env = dict(os.environ)
    env.setdefault("ASAN_OPTIONS", "detect_leaks=0")
    cmd = [
        CHROME,
        "--no-sandbox",
        "--no-first-run",
        "--no-default-browser-check",
        "--disable-gpu",
        "--user-data-dir=" + UDD,
        "--profile-directory=" + profile_dir,
        # Don't pop a window; we just need profile init + Local State write.
        "--no-startup-window",
        "--enable-logging=stderr",
        "--v=0",
    ]
    sys.stderr.write(f"[setup] launching Chrome to init profile '{profile_dir}'\n")
    p = subprocess.Popen(
        cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
    )
    # Give Chrome time to create the profile directory and write Local State.
    deadline = time.time() + 25
    target = os.path.join(UDD, profile_dir, "Preferences")
    while time.time() < deadline:
        if os.path.exists(target):
            time.sleep(2)  # let it settle
            break
        time.sleep(0.5)
    p.terminate()
    try:
        p.wait(timeout=10)
    except subprocess.TimeoutExpired:
        p.kill()
        p.wait()
    # Chrome flushes Local State on shutdown; give it a moment.
    time.sleep(1)


def patch_local_state():
    ls_path = os.path.join(UDD, "Local State")
    with open(ls_path, "r") as f:
        ls = json.load(f)
    info_cache = ls.setdefault("profile", {}).setdefault("info_cache", {})
    if "Profile 1" not in info_cache:
        sys.stderr.write(
            "[setup] WARNING: 'Profile 1' not found in info_cache, creating stub entry\n"
        )
        info_cache["Profile 1"] = {
            "active_time": time.time(),
            "name": "Work (OIDC)",
            "is_omitted": False,
            "is_ephemeral": False,
            "avatar_icon": "chrome://theme/IDR_PROFILE_AVATAR_26",
            "background_apps": False,
        }
    entry = info_cache["Profile 1"]
    # These are exactly the keys ManagedProfileCreator/OidcManagedProfileCreationDelegate
    # set when a real OIDC enrollment completes.
    entry["profile_management_oidc_auth_token"] = "stored_auth_token"
    entry["profile_management_oidc_id_token"] = "stored_id_token"
    entry["profile_management_oidc_state"] = ""
    entry["profile_management_id"] = PROFILE_MGMT_ID
    entry["dasherless_management"] = False
    entry["name"] = "Work (OIDC victim)"
    entry["is_omitted"] = False
    with open(ls_path, "w") as f:
        json.dump(ls, f)
    sys.stderr.write(
        f"[setup] seeded OIDC identity into Profile 1: {PROFILE_MGMT_ID}\n"
    )


def main():
    if os.path.exists(UDD):
        shutil.rmtree(UDD)
    os.makedirs(UDD)
    run_chrome_once("Default")
    run_chrome_once("Profile 1")
    patch_local_state()
    sys.stderr.write("[setup] done\n")


if __name__ == "__main__":
    main()