import helpers.adb.adb as adb, helpers.fastboot.fastboot as fastboot, install, time, os, requests, urllib.request, sys, helpers.unpackBackup as unpackBackup
from lxml import html

def download_twrp():
    cpu = os.popen('adb shell getprop ro.product.board').read()
    cpu = cpu.rstrip('\n')
    r = urllib.request.urlopen('https://dl.twrp.me/'+str(cpu))
    text = r.read().decode('utf-8')
    tree = html.fromstring(text)
    urls = tree.xpath('//a/@href')
    downloads = []
    for i in urls:
        if "img" in i:
            downloads.append(i)
    url_to_download = "https://dl.twrp.me"+downloads[0]
    url_to_download = url_to_download.replace('.html', '')
    s = requests.Session()
    s.headers.update({'referer': url_to_download})
    img = s.get(url_to_download)
    with open("twrp.img", 'wb') as f:
        f.write(img.content)

def device_pull():
    install.main()
    print("connecting device...")
    try:
        device = adb.connect_device()
    except:
        print("Unable to connect to device. Ensure that Debug Mode is enabled and USB Debugging is allowed")
        sys.exit()
    print("getting twrp...")
    download_twrp()
    adb.reboot_bootloader()
    fastboot.devices()
    fastboot.unlock()
    fastboot.boot('./twrp.img')
    time.sleep(20)
    try:
        device = adb.connect_device()
    except:
        time.sleep(10)
        device = adb.connect_device()
    command = adb.shell(device, "whoami")
    print(command)

def menu():
    print('Acquisition Script for Mobile Resources')
    print("---------------------------------------")
    print("[1] Device Pull")
    print("[2] Application Pull")
    print("[3] ADB Backup")
    print("[99] Exit")

def main():
    menu()
    choice = int(input("Please select an option: "))
    match choice:
        case 1:
            device_pull()
        case 2:
            pass
        case 3: 
            adb.backup()
            os.system("python helpers/unpackBackup.py -i ./backup.ab")
        case 99:
           sys.exit() 

if __name__ == '__main__':
    main()