#!/bin/python3
import base64
import os
import sys
import argparse
import logging
import json
import requests
from cryptography.hazmat.primitives.serialization import load_pem_private_key
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import hashes, serialization

from lib.diffie_hellman import *
from subject import main

#  Identity attributes
# {'username' : '', 'full_name' : '', 'email' : '', public_key : '' }

logging.basicConfig(format='%(levelname)s\t- %(message)s')
logger = logging.getLogger()
logger.setLevel(logging.INFO)

state = main(sys.argv)

BASE_DIR = os.path.join(os.path.expanduser('~'), '.sio/')

# org - username - password - credentials file - session file
def createSession(args):

    parser = argparse.ArgumentParser()

    parser.add_argument("-k", '--key', nargs=1, help="Path to the key file")
    parser.add_argument("-r", '--repo', nargs=1, help="Address:Port of the repository")
    parser.add_argument("-v", '--verbose', help="Increase verbosity", action="store_true")

    parser.add_argument('org', nargs='?', default=None)
    parser.add_argument('username', nargs='?', default=None)
    parser.add_argument('password', nargs='?', default=None)
    parser.add_argument('credentials', nargs='?', default=None)
    parser.add_argument('session', nargs='?', default=None)

    args = parser.parse_args()

    if not args.org or not args.username or not args.password or not args.credentials or not args.session:
        logger.error("Need organization, username, credentials and session file")
        sys.exit(1)

    if not os.path.isfile(BASE_DIR + args.credentials):
        logger.error("File '" + args.credentials + "' not found.")
        sys.exit(1)
    
    # perform diffie-hellman key exchange
    # OPTIONAL: ask for parameters
    generator = 2; key_size = 1024
    parameters = generate_parameters(generator, key_size)

    parameters_send = parameters.parameter_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.ParameterFormat.PKCS3).hex()

    # send parameters to server
    try:
        req = requests.post(f'http://{state['REP_ADDRESS']}/user/dh', json=json.dumps({'user': args.username,
                                                                                       'org': args.org,
                                                                                       'parameters': parameters_send}))
        req.raise_for_status()

    except requests.exceptions.HTTPError:
        logger.error("%d: %s", req.status_code, req.json()['error'])
        sys.exit(-1)

    except requests.exceptions.RequestException as errex:
        logger.error("Failed to obtain response from server")
        sys.exit(-1)

    private_key, public_key = generate_key_pair(parameters)

    response = req.json()
    server_public_key = serialization.load_pem_public_key(bytes.fromhex(response['public_key']))
    derived_key = derive_keys(private_key, server_public_key)

    # save derived key
    with open(BASE_DIR + 'derived_key.pem', 'wb') as f:
        f.write(derived_key)

    # send public key to server
    public_key_content = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    ).hex()

    try:
        req = requests.post(f'http://{state['REP_ADDRESS']}/user/dh', json=json.dumps({'user': args.username,
                                                                                       'org': args.org,
                                                                                       'public_key': public_key_content}))
        req.raise_for_status()

    except requests.exceptions.HTTPError:
        logger.error("%d: %s", req.status_code, req.json()['error'])
        sys.exit(-1)

    except requests.exceptions.RequestException as errex:
        logger.error("Failed to obtain response from server")
        sys.exit(-1)

    if req.status_code == 200:
        logger.info("Diffie-Hellman key exchange completed successfully")
    else:
        logger.error("Failed to complete Diffie-Hellman key exchange")
        sys.exit(-1)

    try:
        session = {'org': args.org, 'username': args.username}
        req = requests.post(f'http://{state['REP_ADDRESS']}/user/login', json=json.dumps(session))
        req.raise_for_status()

    except requests.exceptions.HTTPError:
        logger.error("%d: %s", req.status_code, req.json()['error'])
        sys.exit(-1)

    except requests.exceptions.RequestException as errex:
        logger.error("Failed to obtain response from server")
        sys.exit(-1)

    req = decrypt(bytes.fromhex(req.text), derived_key).decode('utf-8')
    response = json.loads(req)
    challenge = response['challenge']

    with open(BASE_DIR + args.credentials, 'rb') as f:
        try:
            key = load_pem_private_key(f.read(), password=args.password.encode("utf-8") if args.password else None)
        except ValueError:
            logger.error("Invalid password")
            sys.exit(-1)

    signature = key.sign(
        challenge.encode('utf-8'),
        padding.PKCS1v15(),
        hashes.SHA256()
    )

    try:
        json_payload = json.dumps({'signature' : base64.b64encode(signature).decode('utf-8')})
        json_payload = encrypt(json_payload, derived_key).hex()

        headers = {
            'Authorization': response['token'],
            'Content-Type': 'application/octet-stream'
        }

        req = requests.post(f'http://{state['REP_ADDRESS']}/user/login', data=json_payload, headers=headers)
        req.raise_for_status()

    except requests.exceptions.HTTPError:
        logger.error("%d: %s", req.status_code, req.json()['error'])
        sys.exit(-1)

    except requests.exceptions.RequestException as errex:
        logger.error("Failed to obtain response from server")
        sys.exit(-1)

    if req.status_code == 201:
        try:
            response_data = json.loads(decrypt(bytes.fromhex(req.text), derived_key).decode('utf-8'))
        except Exception:
            logger.error("Failed to decrypt the content")
            sys.exit(1)
        session = response_data | {'derived_key': derived_key.hex()}
        with open(BASE_DIR + args.session, 'w') as f:
            json.dump(session, f)
        logger.info("Session created successfully")
        sys.exit(0)
    else:
        logger.error("Failed to create session")
        sys.exit(-1)

if __name__ == '__main__':
    createSession(sys.argv[1:])