#!/usr/bin/env python3

import logging
import socket
import os


def chunk(s, pos, l):
    '''
    '''

    lev = 3
    retries = 3

    datalist = [
        int.to_bytes(lev, 1, "little"),   # level
        int.to_bytes(pos, 2, "little"),   # position
        l * b'\x00'
    ]

    packet = b''
    packet += b'\xf0\x00' # static key
    for data in datalist:
       packet += int.to_bytes(len(data), 2, "big")
       packet += data

    data = None
    for i in range(0, retries):
        try:
            s.send(packet)
            data = s.recv(4096)
        except Exception as e:
            if i < retries - 1:
                logging.warning(f"{s}: {e}")
            else:
                logging.fatal(f"{s}: {e}, exiting")
        else:
            logging.debug(f"public-key {l}-bytes chunk {pos}: received")
            return data

    exit(1)

logging.basicConfig(level='DEBUG', format='%(asctime)s %(levelname)s %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
logging.info('start')

s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.settimeout(5)
s.connect(('127.0.0.1', 584))

# receive public key from keyserver
pk_keyserver = b''
for pos in range(0, 910):
    if pos == 909:
        data = chunk(s, pos, 183)
    else:
        data = chunk(s, pos, 1152)
    pk_keyserver += data[11:]

# read public key from filesystem
pk_filesystem = open('/etc/pqconnect/keys/mceliece_pk', 'rb').read()
pk_filesystem += open('/etc/pqconnect/keys/x25519_pk', 'rb').read()

# check
assert(pk_keyserver == pk_filesystem)
logging.info(f"public-key from keyserver received OK")
exit(0)
