#!/usr/bin/env python3
import subprocess
import os
import time
import threading
import shutil

BUNDLE_ID = "com.brave.Browser"
# BUNDLE_ID = "com.google.chrome"
MANAGED_PLIST = f"/Library/Managed Preferences/{BUNDLE_ID}.plist"

# === Brave-specific feature lockdown ===
managed_policies = {
    # "RestoreOnStartup": 5,                            # Open new tab on startup
    # "HomepageIsNewTabPage": True,                     # Use New Tab as homepage
    "AutofillAddressEnabled": False,                  # Disable address autofill
    "AutofillCreditCardEnabled": False,               # Disable credit card autofill
    "BackgroundModeEnabled": False,                   # Disable background tasks
    "BookmarkBarEnabled": False,                      # Disable bookmark bar
    "BraveAIChatEnabled": False,                      # Disable Brave AI features
    "BraveNewsDisabled": True,                        # Disable Brave News
    "BravePlaylistEnabled": False,                    # Disable Brave Playlist
    "BraveRewardsDisabled": True,                     # Disable Brave Rewards
    "BraveSpeedreaderEnabled": False,                 # Disable Speedreader
    "BraveStatsPingEnabled": False,                   # Disable telemetry pings
    "BraveTalkDisabled": True,                        # Disable Brave Talk
    "BraveVPNDisabled": True,                         # Disable Brave VPN
    "BraveWalletDisabled": True,                      # Disable Brave Wallet
    "BraveWalletShowInToolbar": False,                # Hide wallet icon
    "BraveWebDiscoveryEnabled": False,                # Disable web discovery
    "BrowserSignin": 0,                               # Disable browser sign-in
    "DefaultGeolocationSetting": 2,                   # Block location access
    "EnableMediaRouter": False,                       # Disable Google Cast
    "GoogleSearchSidePanelEnabled": False,            # Disable Google Search side panel
    "PasswordLeakDetectionEnabled": False,            # Disable password leak check
    "PasswordManagerEnabled": False,                  # Disable password manager
    "PasswordManagerPasskeysEnabled": False,          # Disable passkeys
    "PasswordSharingEnabled": False,                  # Disable password sharing
    "PaymentMethodQueryEnabled": False,               # Disable payment autofill
    "PrivacySandboxAdMeasurementEnabled": False,      # Disable Privacy Sandbox ads
    "PrivacySandboxAdTopicsEnabled": False,           # Disable ad topics
    "PrivacySandboxPromptEnabled": False,             # Disable Privacy Sandbox prompt
    "PrivacySandboxSiteEnabledAdsEnabled": False,     # Disable site-enabled ads
    "PromotionsEnabled": False,                       # Disable promotional content
    "ShowHomeButton": False,                          # Hide Home button
    "SideSearchEnabled": False,                       # Disable side search
    "SpellcheckEnabled": False,                       # Disable spellcheck
    "SuggestedContentEnabled": False,                 # Disable content suggestions
    "SyncDisabled": True,                             # Disable Brave Sync
    "TorDisabled": True,                              # Disable Tor integration
    "TranslateEnabled": False,                        # Disable translate feature
    "UserFeedbackAllowed": False,                     # Disable sending feedback
    "DeviceMetricsReportingEnabled": False,
    # "MetricsReportingEnabled": False,
    "SafeBrowsingExtendedReportingEnabled": False,
    "DefaultBrowserSettingEnabled": False,
  "AlternateErrorPagesEnabled": False,
  "URLKeyedAnonymizedDataCollectionEnabled": False,
  "SpellCheckServiceEnabled": False,
  "SearchSuggestEnabled": False,
    "SigninInterceptionEnabled": False
}

def run(cmd, use_sudo=False, silent=False):
    if use_sudo:
        cmd = ["sudo"] + cmd
    if not silent:
        print("→", " ".join(cmd))
    subprocess.run(cmd, check=True)

def keep_sudo_alive():
    """Refresh sudo timestamp every 60 s."""
    while True:
        subprocess.run(["sudo", "-v"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        time.sleep(60)

def ensure_managed_dir():
    os.makedirs("/Library/Managed Preferences", exist_ok=True)
    run(["chown", "root:wheel", "/Library/Managed Preferences"], use_sudo=True)
    run(["chmod", "755", "/Library/Managed Preferences"], use_sudo=True)

def set_policy(key, value):
    # Remove existing key before adding (avoid duplicate errors)
    subprocess.run(
        ["sudo", "/usr/libexec/PlistBuddy", "-c", f"Delete :{key}", MANAGED_PLIST],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
    )
    # Choose correct PlistBuddy type
    if isinstance(value, bool):
        cmd = ["Add", f":{key}", "bool", "true" if value else "false"]
    elif isinstance(value, int):
        cmd = ["Add", f":{key}", "integer", str(value)]
    elif isinstance(value, str):
        cmd = ["Add", f":{key}", "string", value]
    elif isinstance(value, list):
        # array
        run(["/usr/libexec/PlistBuddy", "-c", f"Add :{key} array", MANAGED_PLIST], use_sudo=True)
        for v in value:
            run(["/usr/libexec/PlistBuddy", "-c", f"Add :{key}: string {v}", MANAGED_PLIST], use_sudo=True)
        return
    else:
        print(f"⚠️ Unsupported type for key {key}")
        return
    run(["/usr/libexec/PlistBuddy", "-c", " ".join(cmd), MANAGED_PLIST], use_sudo=True)

def set_managed_policies():
    ensure_managed_dir()
    for key, value in managed_policies.items():
        set_policy(key, value)

def refresh_preferences():
    run(["killall", "cfprefsd"], use_sudo=True)

def cleanup_brave_data():
    """Optionally remove crypto-related folders."""
    profile = os.path.expanduser("~/Library/Application Support/BraveSoftware/Brave-Browser/Default")
    for folder in ["brave_rewards", "ethereum_wallet", "brave_vpn", "ai_chat"]:
        path = os.path.join(profile, folder)
        if os.path.exists(path):
            print(f"🧹 Removing {path}")
            shutil.rmtree(path, ignore_errors=True)

if __name__ == "__main__":
    print("🔐 Applying full Brave Browser hardening policies (mandatory).")
    subprocess.run(["sudo", "-v"], check=True)  # ask for password once
    threading.Thread(target=keep_sudo_alive, daemon=True).start()

    set_managed_policies()
    cleanup_brave_data()
    refresh_preferences()

    print(f"\n✅ All mandatory policies applied under:\n   {MANAGED_PLIST}")
    print("   Restart Brave and check brave://policy — all should show Level = Mandatory.")
    print("   You can re-run this script anytime to refresh or update settings.")