#!/usr/bin/env python3

import argparse
import base64
import configparser
import json
import os
import pathlib
import subprocess
import sys
import traceback
import urllib.parse

# use requests instead of urllib.request for keep-alive connection
import requests

def sign(data, keyfile):
    p = subprocess.run(
        ['openssl', 'pkeyutl', '-sign', '-inkey', keyfile, '-pkeyopt', 'digest:sha256'],
        input=base64.b64decode(data),
        capture_output=True)
    return base64.b64encode(p.stdout).decode()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Fake the MargTools application.')
    parser.add_argument('url', type=urllib.parse.urlparse, help='bc-digsign:// url')
    parser.add_argument('-k', '--user-key', type=pathlib.Path, help='key file')
    parser.add_argument('-c', '--user-cert', type=pathlib.Path, help='certificate file')
    args = parser.parse_args()

    try:
        # parse query string
        params = urllib.parse.parse_qs(args.url.query)
        url = params['baseUrl'][0]
        token = params['accessToken'][0]

        # if missing, get user key and cert from section [url] in ~/.margfools
        config = configparser.ConfigParser()
        config.read(os.path.expanduser('~') + '/.margfools')
        if not args.user_key:
            args.user_key = config.get(url, 'user-key')
        if not args.user_cert:
            args.user_cert = config.get(url, 'user-cert')
        if not args.user_key or not args.user_cert:
            print('user key and/or certificate not specified', file=sys.stderr)
            sys.exit(1)

        user_keyfile = args.user_key
        user_cert = ''.join(line.strip() for line in open(args.user_cert) if not line.startswith('-----'))

        session = requests.Session()
        headers={'Authorization': f'Bearer {token}'}

        # delete old signing session
        r = session.delete(f'{url}/signatures/{params["startSigningToken"][0]}', headers=headers)

        # register a certificate or sign a document, makes no difference to us
        if params.get('registerCertificate'):
            q = {'registerCertificate': 1}
        else:
            q = {'documentId': params['documentId'][0]}
        r = session.post(f'{url}/signatures?{urllib.parse.urlencode(q)}', headers=headers)

        # get signature request and mix in my secrets and publics
        request = json.loads(r.text)
        request['AuthenticationToken'] = token
        request['CertificatePublicKey'] = user_cert

        # keep signing whatever they send us
        while True:
            for name in ('AttachmentHashes', 'XmlHashes'):
                if request.get(name) is not None:
                    request[f'Signed{name}'] = [sign(e, user_keyfile) for e in request[name]]

            r = session.put(f'{url}signatures/{request["SignatureRequestId"]}',
                headers=headers | {'Content-Type': 'application/json; charset=utf-8'},
                data=json.dumps(request).encode())
            if not r.text:
                break
            request |= json.loads(r.text)

    except:
        traceback.print_exc()