#!/usr/bin/env vpython3
# coding: utf-8

import os
import sys
import binascii
import datetime
import PyKCS11

from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa

from asn1crypto import x509 as asn1x509
from asn1crypto import keys as asn1keys
from asn1crypto import pem as asn1pem
from asn1crypto import util as asn1util

class HSM:
    def __init__(self, dllpath):
        self.pkcs11 = PyKCS11.PyKCS11Lib()
        self.pkcs11.load(dllpath)
        self.session = None

    def getSlot(self, label):
        slots = self.pkcs11.getSlotList(tokenPresent=True)
        for slot in slots:
            info = self.pkcs11.getTokenInfo(slot)
            try:
                if info.label.split('\0')[0].strip() == label:
                    return slot
            except AttributeError:
                continue
        return None

    def create(self, label, pin, sopin):
        slot = self.getSlot(label)
        if slot is not None:
            return
        slot = self.pkcs11.getSlotList(tokenPresent=True)[-1]
        self.pkcs11.initToken(slot, sopin, label)
        session = self.pkcs11.openSession(slot, PyKCS11.CKF_SERIAL_SESSION | PyKCS11.CKF_RW_SESSION)
        session.login(sopin, user_type=PyKCS11.CKU_SO)
        session.initPin(pin)
        session.logout()
        session.closeSession()

    def login(self, label, pin):
        slot = self.getSlot(label)
        if slot is None:
            return
        self.session = self.pkcs11.openSession(slot, PyKCS11.CKF_SERIAL_SESSION | PyKCS11.CKF_RW_SESSION)
        self.session.login(pin)

    def logout(self):
        if self.session is not None:
            self.session.logout()
            self.session.closeSession()
            self.session = None

    def gen_privkey(self, label, key_id, key_length=2048):
        # label - just a label for identifying objects
        # key_id has to be the same for both objects, it will also be necessary
        #     when importing the certificate, to ensure it is linked with these keys.
        # key_length - key-length in bits

        public_template = [
            (PyKCS11.CKA_CLASS, PyKCS11.CKO_PUBLIC_KEY),
            (PyKCS11.CKA_TOKEN, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_PRIVATE, PyKCS11.CK_FALSE),
            (PyKCS11.CKA_MODULUS_BITS, key_length),
#            (PyKCS11.CKA_PUBLIC_EXPONENT, (0x01, 0x00, 0x01)),
            (PyKCS11.CKA_ENCRYPT, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_VERIFY, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_VERIFY_RECOVER, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_WRAP, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_LABEL, label),
            (PyKCS11.CKA_ID, key_id)
#            (PyKCS11.CKA_KEY_TYPE, PyKCS11.CKK_RSA),
#            (PyKCS11.CKA_SENSITIVE, PyKCS11.CK_FALSE),
        ]

        private_template = [
            (PyKCS11.CKA_CLASS, PyKCS11.CKO_PRIVATE_KEY),
            (PyKCS11.CKA_TOKEN, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_PRIVATE, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_DECRYPT, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_SIGN, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_SIGN_RECOVER, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_UNWRAP, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_LABEL, label),
            (PyKCS11.CKA_ID, key_id)
#            (PyKCS11.CKA_SENSITIVE, PyKCS11.CK_TRUE),
        ]

        self.session.generateKeyPair(public_template, private_template)

    def cert_save(self, cert, label, subject, key_id):
        cert_template = [
            (PyKCS11.CKA_CLASS, PyKCS11.CKO_CERTIFICATE),
            (PyKCS11.CKA_CERTIFICATE_TYPE, PyKCS11.CKC_X_509),
            (PyKCS11.CKA_TOKEN, PyKCS11.CK_TRUE),
            (PyKCS11.CKA_LABEL, label.encode('utf-8')),
            (PyKCS11.CKA_ID, key_id),  # must be set, and DER see Table 24, X.509 Certificate Object Attributes
            (PyKCS11.CKA_SUBJECT, subject.encode('utf-8')),  # must be set and DER, see Table 24, X.509 Certificate Object Attributes

            #(PyKCS11.CKA_PRIVATE, PyKCS11.CK_FALSE),
            #(PyKCS11.CKA_TRUSTED, PyKCS11.CK_TRUE),
            #(PyKCS11.CKA_SENSITIVE, PyKCS11.CK_FALSE),
            #(PyKCS11.CKA_ENCRYPT, PyKCS11.CK_TRUE),
            #(PyKCS11.CKA_VERIFY, PyKCS11.CK_TRUE),
            #(PyKCS11.CKA_MODIFIABLE, PyKCS11.CK_TRUE),
#            (PyKCS11.CKA_ISSUER, cert.Issuer);
#            (PyKCS11.CKA_SERIAL_NUMBER,cert.SerialNumber)
            (PyKCS11.CKA_VALUE, cert),  # must be BER-encoded

        ]

        self.session.createObject(cert_template)

    def cert_load(self, keyID):
        rec = self.session.findObjects([(PyKCS11.CKA_CLASS, PyKCS11.CKO_CERTIFICATE), (PyKCS11.CKA_ID, keyID)])
        if len(rec) == 0:
            return None
        value = bytes(rec[0].to_dict()['CKA_VALUE'])
        return value

    def certsign(self, sn, pubKey, subject, until, caprivKey):
        tbs = asn1x509.TbsCertificate({
            'version': 'v1',
            'serial_number': sn,
            'issuer': asn1x509.Name.build({
                'common_name': 'hsm CA',
            }),
            'subject': asn1x509.Name.build({
                'common_name': subject,
            }),
            'signature': {
                'algorithm': 'sha256_rsa',
                'parameters': None,
            },
            'validity': {
                'not_before': asn1x509.Time({
                    'utc_time': datetime.datetime.now(tz=asn1util.timezone.utc) - datetime.timedelta(days=1),
                }),
                'not_after':  asn1x509.Time({
                    'utc_time': until,
                }),
            },
            'subject_public_key_info': {
                'algorithm': {
                    'algorithm': 'rsa',
                    'parameters': None,
                },
                'public_key': pubKey
            }
        })

        # Sign the TBS Certificate
        data = tbs.dump()
        value = self.session.sign(caprivKey, data, PyKCS11.Mechanism(PyKCS11.CKM_SHA256_RSA_PKCS, None))
        value = bytes(bytearray(value))

        cert = asn1x509.Certificate({
            'tbs_certificate': tbs,
            'signature_algorithm': {
                'algorithm': 'sha256_rsa',
                'parameters': None,
            },
            'signature_value': value,
        })
        return cert.dump()

    def ca_gen(self, label, keyID, subject):
        privKey = self.session.findObjects([(PyKCS11.CKA_CLASS, PyKCS11.CKO_PRIVATE_KEY), (PyKCS11.CKA_ID, keyID)])[0]
        pubKey = self.session.findObjects([(PyKCS11.CKA_CLASS, PyKCS11.CKO_PUBLIC_KEY), (PyKCS11.CKA_ID, keyID)])[0]

        modulus = self.session.getAttributeValue(pubKey, [PyKCS11.CKA_MODULUS])[0]
        modulus = binascii.hexlify(bytearray(modulus)).decode("utf-8")
        exponent = self.session.getAttributeValue(pubKey, [PyKCS11.CKA_PUBLIC_EXPONENT])[0]
        exponent = binascii.hexlify(bytearray(exponent)).decode("utf-8")
        pubKey = asn1keys.RSAPublicKey({
            'modulus':int('0x'+modulus, 16),
            'public_exponent':int('0x'+exponent, 16)
        })
        #pubKey = asn1keys.RSAPublicKey.load(pubKey.dump())
        until = datetime.datetime.now(tz=asn1util.timezone.utc) + datetime.timedelta(days=365*10)
        der_bytes = self.certsign(1, pubKey, subject, until, privKey)
        self.cert_save(der_bytes, label, subject, keyID)

    def ca_sign(self, keyID, label, sn, subject, days, cakeyID):
        caprivKey = self.session.findObjects([(PyKCS11.CKA_CLASS, PyKCS11.CKO_PRIVATE_KEY), (PyKCS11.CKA_ID, cakeyID)])[0]

        pubKey = self.session.findObjects([(PyKCS11.CKA_CLASS, PyKCS11.CKO_PUBLIC_KEY), (PyKCS11.CKA_ID, keyID)])[0]
        modulus = self.session.getAttributeValue(pubKey, [PyKCS11.CKA_MODULUS])[0]
        modulus = binascii.hexlify(bytearray(modulus)).decode("utf-8")
        exponent = self.session.getAttributeValue(pubKey, [PyKCS11.CKA_PUBLIC_EXPONENT])[0]
        exponent = binascii.hexlify(bytearray(exponent)).decode("utf-8")
        pubKey = asn1keys.RSAPublicKey({
            'modulus':int('0x'+modulus, 16),
            'public_exponent':int('0x'+exponent, 16)
        })
        #pubKey = asn1keys.RSAPublicKey.load(pubKey.dump())
        until = datetime.datetime.now(tz=asn1util.timezone.utc) + datetime.timedelta(days=days)
        der_bytes = self.certsign(sn, pubKey, subject, until, caprivKey)
        self.cert_save(der_bytes, label, subject, keyID)

    def cert_export(self, fname, keyID):
        der_bytes = self.cert_load(keyID)
        pem_bytes = asn1pem.armor('CERTIFICATE', der_bytes)
        open(fname+'.der', 'wb').write(der_bytes)
        open(fname+'.pem', 'wb').write(pem_bytes)
