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

filename = ''

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()
    serialno = adb.serialno()
    logger.add_log(filename, "Serial Number: "+str(serialno))
    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")
        logger.add_log(filename, "Unable to connect to device. Ensure that Debug Mode is enabled and USB Debugging is allowed")
        sys.exit()
    serialno = adb.serialno()
    logger.add_log(filename, "Serial Number: "+str(serialno))
    print("pulling data...")
    adb.pull(device, "/sdcard", "./sdcard")
    os.system('zip -r logical.zip ./sdcard')
    os.system("rm -r ./sdcard")
    hash = os.popen("md5sum logical.zip").read()
    logger.add_log(filename, "MD5 Hash: "+str(hash))

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: "))
    global filename
    match choice:
        case 1:
            filename = logger.start_logger()
            logger.add_log(filename, "Type: Physical Image")
            device_pull()
            logger.stop_logger(filename)
        case 2:
            filename = logger.start_logger()
            logger.add_log(filename, "Type: Logical Image")
            logicpull()
            logger.stop_logger(filename)
        case 3:
            filename = logger.start_logger()
            logger.add_log(filename, "Type: ADB Backup")
            serialno = adb.serialno()
            logger.add_log(filename, "Serial Number: " + str(serialno))
            adb.backup()
            input("After confirming the backup operation, press enter to continue.")
            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")
            logger.stop_logger(filename)
        case 99:
           sys.exit() 

if __name__ == '__main__':
    main()