#!/usr/bin/env python3

import binascii
import argparse
import logging
import socket
import os

def query(s, pos, lev, padding, retries):
    '''
    '''

    logging.debug('query: pos = %d, lev = %d, padding = %d, retries = %d' % (pos, lev, padding, retries))

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

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

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

if __name__ == '__main__':

    # argparser
    parser = argparse.ArgumentParser(description='keyserver-query.py')
    parser.add_argument('-v', '--verbose',
                        action='count',
                        help='verbosity',
                        default=0)
    parser.add_argument('-l', '--level',
                        action='store',
                        help='level',
                        type=int,
                        default=0)
    parser.add_argument('-p', '--position',
                        action='store',
                        help='position',
                        type=int,
                        default=0)
    parser.add_argument('-t', '--timeout',
                        action='store',
                        help='timeout',
                        type=int,
                        default=1)
    parser.add_argument('-r', '--retries',
                        action='store',
                        help='retries',
                        type=int,
                        default=1)
    parser.add_argument('IP', type=str, help='Keyserver IP address')
    parser.add_argument('PORT', type=int, help='Keyserver port')
    args = parser.parse_args()

    ip = args.IP
    port = args.PORT
    level = args.level
    if (level < 0 or level > 4):
        parser.error('-l/--level must be a number <0-3>')
    position = args.position
    if (position < 0):
        parser.error('-p/--position must be a number <0-%s>' % (positions[level]))
    retries = args.retries

    timeout = args.timeout
        
    # verbosity
    LOG_LEVELS = ["INFO", "DEBUG"]
    if (args.verbose > len(LOG_LEVELS) - 1):
        args.verbose = len(LOG_LEVELS) - 1
    logging.basicConfig(level=LOG_LEVELS[args.verbose])

    # start
    logging.debug('IP = %s' % (ip))
    logging.debug('PORT = %d' % (port))
    logging.debug('level = %d' % (level))
    logging.debug('position = %d' % (position))

    logging.debug('UDP socket')
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    logging.debug('UDP socket created')
    s.settimeout(timeout)
    logging.debug('UDP socket timeout set to %d' % (timeout))
    s.connect((ip, port))
    logging.debug('UDP socket connected to %s:%d' % (ip, port))

    # padding
    padding = -1
    if level == 0:
        padding = 32
    if level == 1:
        padding = 832
    if level == 2:
        padding = 1152
        if position == 25:
            padding = 320
    if level == 3:
        padding = 1152
        if position == 909:
            padding = 183

    # query
    data = query(s, position, level, padding, retries)
    if not data:
        exit(1)

    print(binascii.hexlify(data).decode('ascii'))
    exit(0)

