#!/usr/bin/env python3

import sys
import socket
import struct
import hashlib
import os
import time

from getpass import getpass

# update types
FLASH  = 0
SPIFFS = 1

def readMessage(sock):
    data = sock.recv(4)
    length, = struct.unpack(">I", data)
    data = sock.recv(1)
    success, = struct.unpack("B", data)
    data = sock.recv(length-1)
    message = data.decode('utf-8')

    return success, message

# read the salt from the config file
with open("../data/etc/auth", "r") as authFile:
    lineno = 0
    for line in authFile:
        if lineno == 1:
            SALT = line.strip()
            break
        lineno += 1

_, host, port, updatetypestr, filename = sys.argv

if updatetypestr == "flash":
    updatetype = FLASH
elif updatetypestr == "spiffs":
    updatetype = SPIFFS
else:
    print(f"Invalid update type {updatetypestr}. Valid types are: flash, spiffs")
    exit(1)

# read and store the password from the user
pwd = getpass()

s = socket.create_connection( (host, int(port)) )

data = s.recv(4)
length, = struct.unpack(">I", data)

if length != 4:
    print(f"Unexpected challenge length: {length}")
    exit(1)

data = s.recv(4)
challenge, = struct.unpack(">I", data)

print(f"Challenge: {challenge}")

# build response string
responsestr = pwd + ":" + str(challenge) + ":" + SALT

m = hashlib.sha256()
m.update(responsestr.encode('utf-8'))
response = m.hexdigest()

print(f"Response: {response}")

data = struct.pack(">I64s", len(response), response.encode('ascii') )
s.send(data)

success, message = readMessage(s)
print(f"Server message: {message}")

if not success:
    print("Failed.")
    exit(1)

print() # for proper progress display

with open(filename, "rb") as binfile:
    filesize = os.stat(filename).st_size
    messagesize = filesize + 1 # for type byte

    data = struct.pack(">I", messagesize)
    s.send(data)

    data = struct.pack("B", updatetype)
    s.send(data)

    sent_bytes = 0
    while True:
        data = binfile.read(1024)
        if not data:
            break
        s.send(data)
        sent_bytes += len(data)

        print(f"\rSent {sent_bytes} of {filesize} bytes.")

success, message = readMessage(s)
print(f"Server message: {message}")

if not success:
    print("Failed.")
    exit(1)