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...")
    if os.path.isfile('./twrp.img'):
        pass
    else:
        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()
    adb.pull(device, "/dev/block/sda", "./sda.img")

def logicpull():
    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("pulling data...")
    adb.pull(device, "/sdcard", "./sdcard")
    os.system('zip -r logical.zip ./sdcard')
    os.system("rm -r ./sdcard")

def menu():
    print('Acquisition Script for Mobile Resources')
    print("---------------------------------------")
    print("[1] Physical Image")
    print("[2] Logical Image")
    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:
            logicpull()
        case 3: 
            adb.backup()
            os.system("dd if=backup.ab bs=1 skip=24 | python3 -c 'import zlib,sys;sys.stdout.buffer.write(zlib.decompress(sys.stdin.buffer.read()))' > backup.tar")
        case 99:
           sys.exit() 

if __name__ == '__main__':
    main()