import typing
import subprocess
import hashlib
import re
import os
import OpenSSL.crypto as cr


class HashAlg:
    __slots__ = ('len', 'hashlib_name', 'mech', 'ident')

    def __init__(self, len, hashlib_name, mech, ident):
        self.len = len
        self.hashlib_name = hashlib_name
        self.mech = mech
        self.ident = ident

HASH_ALG = {
    'SHA1': HashAlg(20, 'sha1', 'SHA1-RSA-PKCS', b'\x30\x21\x30\x09\x06\x05\x2b\x0e\x03\x02\x1a\x05\x00\x04\x14'),
    'SHA256': HashAlg(32, 'sha256', 'SHA256-RSA-PKCS', b'\x30\x31\x30\x0d\x06\x09\x60\x86\x48\x01\x65\x03\x04\x02\x01\x05\x00\x04\x20'),
    'SHA384': HashAlg(48, 'sha384', 'SHA384-RSA-PKCS', b'\x30\x41\x30\x0d\x06\x09\x60\x86\x48\x01\x65\x03\x04\x02\x02\x05\x00\x04\x30'),
    'SHA512': HashAlg(64, 'sha512', 'SHA512-RSA-PKCS', b'\x30\x51\x30\x0d\x06\x09\x60\x86\x48\x01\x65\x03\x04\x02\x03\x05\x00\x04\x40'),
}


class Key:
    __slots__ = ('reader', 'id', 'cert_data', 'cert', 'use_rsa_pkcs', 'pin')

    reader: str
    id: str
    cert_data: bytes
    cert: typing.Any
    use_rsa_pkcs: bool
    pin: str

    def __init__(self, reader: str, id: str) -> None:
        self.reader = reader
        self.id = id
        cert_data = self.cert_data = subprocess.check_output(['pkcs11-tool', '--slot', reader, '-y', 'cert', '-r', '-d', id])
        self.cert = cr.load_certificate(cr.FILETYPE_ASN1, cert_data)
        # TODO: autodetect based on key?
        self.use_rsa_pkcs = True

    def sign(self, msg, hash_alg):
        if self.use_rsa_pkcs:
            mech = 'RSA-PKCS'
            msg = hash_alg.ident + hashlib.new(hash_alg.hashlib_name, msg).digest()
        else:
            mech = hash_alg.mech

        env = os.environ.copy()
        env['PIN'] = self.pin
        proc = subprocess.run(
            ['pkcs11-tool', '--slot', self.reader, '--id', self.id, '-m', mech, '--pin', 'env:PIN', '--sign'],
            input = msg,
            stdout = subprocess.PIPE,
            env = env,
            check = True,
        )

        return proc.stdout

class Store:
    __slots__ = ('_impl',)

    def __init__(self):
        self._impl = cr.X509Store()

    def load_file(self, path):
        self._impl.load_locations(path, None)

    def cert_chain(self, key: Key):
        ctx = cr.X509StoreContext(self._impl, key.cert)
        return ctx.get_verified_chain()

def list() -> typing.Iterable[Key]:
    lines = subprocess.check_output(['pkcs11-tool', '--list-slots'], text=True).splitlines()
    readers = []
    for line in lines:
        m = re.fullmatch('Slot ([0-9]+) \\(.*\\): .*', line)
        if m is None:
            continue
        readers.append(m.group(1))

    certs = {}
    for reader in readers:
        lines = subprocess.check_output(['pkcs11-tool', '--slot', reader, '-y', 'cert', '-O'], text=True).splitlines()
        for line in lines:
            m = re.fullmatch('  ID *: *([0-9a-f]+)$', line)
            if m is None:
                continue
            yield Key(reader, m.group(1))