#!/usr/bin/env python3
"""
CryptoVault USB Flasher - Admin Version
R formats USB and flashes bootable cold wallet
Requires admin to format & set up boot sector
"""

import os, sys, time, subprocess, hashlib, shutil, struct
from pathlib import Path

# === COLORS ===
R, G, Y, B, M, C, W, N = '\033[91m', '\033[92m', '\033[93m', '\033[94m', '\033[95m', '\033[96m', '\033[97m', '\033[0m'

def log(msg, color=Y):
    print(f"{color}{msg}{N}")

def get_usb_drives():
    """Detect USB drives on Windows"""
    drives = []
    try:
        result = subprocess.run(['wmic', 'logicaldisk', 'get', 'caption,drivetype,volumename', '/format:csv'], 
                              capture_output=True, text=True, timeout=10)
        for line in result.stdout.strip().split('\n')[1:]:
            parts = line.split(',')
            if len(parts) >= 3 and parts[1] == '2':  # Removable
                drives.append({'letter': parts[2].strip(':'), 'name': parts[3] if len(parts) > 3 else 'USB Drive'})
    except:
        # Fallback
        for letter in 'DEFGHIJKLMNOPQRSTUVWXYZ':
            path = f'{letter}:\\'
            if os.path.exists(path):
                try:
                    import ctypes
                    free = ctypes.c_ulonglong()
                    total = ctypes.c_ulonglong()
                    ctypes.windll.kernel32.GetDiskFreeSpaceExW(path, None, ctypes.pointer(total), ctypes.pointer(free))
                    drives.append({'letter': letter, 'name': 'USB Drive'})
                except: pass
    return drives

def require_admin():
    """Force admin privileges"""
    try:
        import ctypes
        if not ctypes.windll.shell32.IsUserAnAdmin():
            log("🔄 Requesting admin access...", C)
            ctypes.windll.shell32.ShellExecuteW(None, 'runas', sys.executable, ' '.join(sys.argv), None, 1)
            sys.exit(0)
    except: pass

def format_usb(letter):
    """Format USB drive as FAT32"""
    log(f"⚠️  Formatting {letter}:\\ ...", R)
    try:
        # Use Windows format command
        result = subprocess.run(
            ['format', f'{letter}:', '/FS:FAT32', '/V:CryptoVault', '/Q', '/Y'],
            capture_output=True, text=True, timeout=60
        )
        return True
    except Exception as e:
        log(f"Format error: {e}", R)
        return False

def make_bootable(letter):
    """Make USB bootable with MBR"""
    log(f"⚙️  Setting up boot sector on {letter}:\\ ...", C)
    try:
        # Use bootsect to make bootable
        subprocess.run(['bootsect', f'/nt60', f'{letter}:'], capture_output=True, timeout=30)
        return True
    except:
        # Manual MBR setup
        try:
            import ctypes
            disk = f'\\\\.\\{letter}:'
            handle = ctypes.windll.kernel32.CreateFileA(disk, 0x40000000, 0x3, None, 0x3, 0, None)
            if handle != -1:
                # Write boot sector
                boot_code = bytes([0xEB, 0xFE, 0x90] + [0] * 509) + struct.pack('<I', 0xAA55)
                ctypes.windll.kernel32.WriteFile(handle, boot_code, len(boot_code), None, None)
                ctypes.windll.kernel32.CloseHandle(handle)
            return True
        except:
            return False

def copy_files(letter):
    """Copy all wallet files to USB"""
    src = Path(__file__).parent
    dst = Path(f'{letter}:\\')
    
    files = ['index.html', 'START.bat', 'autorun.inf', 'wallet-setup.reg', 'mine.bat', 
             'CHECKSUMS.txt', 'README.txt']
    
    log("📁 Copying wallet files...", C)
    for f in files:
        src_file = src / f
        if src_file.exists():
            shutil.copy2(src_file, dst / f)
            log(f"  ✅ {f}")
    
    # Create directories
    for d in ['keystore', 'wallet-data']:
        (dst / d).mkdir(exist_ok=True)
    
    return True

def write_autounattend(letter):
    """Create autounattend for Windows setup"""
    content = r'''<?xml version="1.0" encoding="utf-8"?>
<autounattend xmlns="urn:schemas-microsoft-com:unattend">
    <settings pass="windowsPE">
        <component name="Microsoft-Windows-International-Core-WinPE">
            <SetupUILanguage><Value>en-US</Value></SetupUILanguage>
        </component>
    </settings>
</autounattend>'''
    Path(f'{letter}:\\autounattend.xml').write_text(content)

def calculate_checksums(letter):
    """Calculate checksums for verification"""
    log("🔒 Calculating checksums...", C)
    chk = {}
    for f in Path(f'{letter}:\\').rglob('*'):
        if f.is_file():
            h = hashlib.sha256()
            h.update(f.read_bytes())
            chk[f.name] = h.hexdigest()[:16]
    Path(f'{letter}:\\CHECKSUMS.txt').write_text('\n'.join(f"{v}  {k}" for k,v in chk.items()))
    return chk

def main():
    os.system('cls' if os.name=='nt' else 'clear')
    require_admin()
    
    print(f"""{C}
╔══════════════════════════════════════════════════════════╗
║   {W}CryptoVault USB Flasher v1.0{N}{C}                        ║
║   {W}Bootable Multi-Coin Cold Wallet Creator{N}{C}              ║
╚══════════════════════════════════════════════════════════╝
{N}""")
    
    # Detect drives
    drives = get_usb_drives()
    if not drives:
        log("❌ No USB drives detected!", R)
        input("\nPress Enter to exit...")
        return
    
    # List drives
    log("📋 Detected USB drives:", W)
    for i, d in enumerate(drives):
        print(f"  [{i+1}] {G}{d['letter']}:\\{N} - {d['name']}")
    print()
    
    # Select drive
    choice = input(f"{Y}Select drive number (1-{len(drives)}): {N}").strip()
    if not choice.isdigit() or int(choice) < 1 or int(choice) > len(drives):
        log("❌ Invalid selection!", R)
        return
    
    drive = drives[int(choice)-1]['letter']
    
    # Confirm
    print(f"\n{R}⚠️  WARNING: This will ERASE all data on {drive}:\\!{N}")
    confirm = input(f"{Y}Type 'YES' to continue: {N}").strip()
    if confirm != 'YES':
        log("❌ Cancelled.", R)
        return
    
    # Flash!
    print(f"\n{G}🔥 Starting flash process...{N}\n")
    
    # Step 1: Format
    if not format_usb(drive):
        log("❌ Format failed!", R)
        return
    
    # Step 2: Make bootable
    if not make_bootable(drive):
        log("⚠️  Boot sector setup failed (continuing anyway)...", Y)
    
    # Step 3: Copy files
    if not copy_files(drive):
        log("❌ File copy failed!", R)
        return
    
    # Step 4: Write autounattend
    write_autounattend(drive)
    
    # Step 5: Checksums
    calculate_checksums(drive)
    
    # Done!
    print(f"""
{G}╔══════════════════════════════════════════════════════════╗
║   ✅ CryptoVault USB is ready!                              ║
║                                                              ║
║   💾 Drive: {drive}:\\                                         ║
║   🔐 Password: Set on first run                              ║
║   ⛏️  Mining: Run mine.bat                                   ║
║   👢 Bootable: Yes (legacy BIOS)                            ║
╚══════════════════════════════════════════════════════════╝
{N}""")
    
    input("\nPress Enter to exit...")

if __name__ == '__main__':
    main()