# Modified from https://github.com/ayushsharma82/ElegantOTA
#
# Requires Python (sudo apt install python3-pip)
#  `python3 -m venv venv` to create the virtual environment
#  `source ./venv/bin/activate` to enter it
#  `pip install -r requirements.txt` to install the libraries
#
# Run using for example:
#   python3 upload_cli.py -i 10.10.10.173 -f ../build/firmware/EMS-ESP-3_7_0-dev_8-ESP32_4M.bin

import argparse
import requests
import hashlib
from urllib.parse import urlparse
import time
from requests_toolbelt import MultipartEncoder, MultipartEncoderMonitor
from tqdm import tqdm
from termcolor import cprint

def print_success(x): return cprint(x, 'green')
def print_fail(x): return cprint(x, 'red')

def upload(file, ip, username, password):

    # Print welcome message
    print()
    print("EMS-ESP Firmware Upload")

    # first check authentication
    emsesp_url = "http://" + f'{ip}'
    parsed_url = urlparse(emsesp_url)
    host_ip = parsed_url.netloc
    
    signon_url = f"{emsesp_url}/rest/signIn"

    signon_headers = {
        'Host': host_ip,
        'User-Agent': 'Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:109.0) Gecko/20100101 Firefox/118.0',
        'Accept': '*/*',
        'Accept-Language': 'de,en-US;q=0.7,en;q=0.3',
        'Accept-Encoding': 'gzip, deflate',
        'Referer': f'{emsesp_url}',
        'Content-Type': 'application/json',
        'Connection': 'keep-alive'
    }

    username_password = {
        "username": username,
        "password": password
    }

    response = requests.post(signon_url, json=username_password, headers=signon_headers, auth=None)
        
    if response.status_code != 200:
        print_fail("Authentication failed (code " + str(response.status_code) + ")")
        return

    print_success("Authentication successful")
    access_token = response.json().get('access_token')

    # start the upload
    with open(file, 'rb') as firmware:
        md5 = hashlib.md5(firmware.read()).hexdigest()
       
        firmware.seek(0)

        encoder = MultipartEncoder(fields={
            'MD5': md5,
            'file': (file, firmware, 'application/octet-stream')}
        )

        bar = tqdm(desc='Upload Progress',
                   total=encoder.len,
                   dynamic_ncols=True,
                   unit='B',
                   unit_scale=True,
                   unit_divisor=1024
                   )

        monitor = MultipartEncoderMonitor(encoder, lambda monitor: bar.update(monitor.bytes_read - bar.n))

        post_headers = {
            'Host': host_ip,
            'User-Agent': 'Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:109.0) Gecko/20100101 Firefox/118.0',
            'Accept': '*/*',
            'Accept-Language': 'de,en-US;q=0.7,en;q=0.3',
            'Accept-Encoding': 'gzip, deflate',
            'Referer': f'{emsesp_url}',
            'Connection': 'keep-alive',
            'Content-Type': monitor.content_type,
            'Content-Length': str(monitor.len),
            'Origin': f'{emsesp_url}',
            'Authorization': 'Bearer ' + f'{access_token}'
        }

        upload_url = f"{emsesp_url}/rest/uploadFile"

        response = requests.post(upload_url, data=monitor, headers=post_headers, auth=None)
        
        bar.close()
        time.sleep(0.1)

        if response.status_code != 200:
            print_fail("Upload failed (code " + response.status.code + ").")
        else:
            print_success("Upload successful.")

        print()

# main
parser = argparse.ArgumentParser(description="EMS-ESP Firmware Upload")
parser.add_argument("-f", "--file", metavar="FILE", required=True, type=str, help="firmware file")
parser.add_argument("-i", "--ip", metavar="IP", type=str, default="ems-esp.local", help="IP address of EMS-ESP")
parser.add_argument("-u", "--username", metavar="USERNAME", type=str, default="admin", help="admin user")
parser.add_argument("-p", "--password", metavar="PASSWORD", type=str, default="admin", help="admin password")
args = parser.parse_args()
upload(**vars(args))