import os
import getpass

from pyasn1.codec.der import decoder as der_decoder

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding, hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.asymmetric import padding as asympad
from cryptography.hazmat.primitives.serialization import load_der_private_key
from cryptography.hazmat.primitives.serialization import load_der_public_key
from cryptography.exceptions import InvalidSignature


NPKI_PATH = "C:\\Program Files\\NPKI\\"
SEED_CBC_SHA1 = "1.2.410.200004.1.15"


def hasher(data):
    s = hashes.Hash(hashes.SHA1(), backend=default_backend())
    s.update(data)
    return s


def _pbkdf1(password, salt, itercnt):
    s = hasher(password + salt)
    dk = s.finalize()
    for i in range(itercnt - 1):
        s = hasher(dk)
        dk = s.finalize()
    return dk


def bit2string(data):
    bits = list(data) + [0] * ((8 - len(data) % 8) % 8)
    blist = []
    for i in range(0, len(bits), 8):
        blist.append(chr(int(''.join(str(b) for b in bits[i:i + 8]), 2)))
    data = ''.join(blist)
    return data


class KoCertificate:
    def __init__(self, pathname, passwd=None):
        self.cert_path = pathname
        self.pub_key = self.get_pubkey()
        if passwd is not None:
            self.pri_key = self.get_prikey(passwd)

    def get_pubkey(self):
        pub_filename = os.path.join(self.cert_path, "signCert.der")

        der_file = open(pub_filename, 'rb').read()
        der_seq, remaining = der_decoder.decode(der_file)
        seq1 = der_seq.getComponentByPosition(0)
        valid = seq1.getComponentByPosition(4)
        self.valid_date = [valid.getComponentByPosition(0),
                           valid.getComponentByPosition(1)]

        seq2 = seq1.getComponentByPosition(5).getComponentByPosition(4)
        self.owner = unicode(str(seq2.getComponentByPosition(0).
                                 getComponentByPosition(1)), 'utf-8')

        # seq3 = seq1.getComponentByPosition(7)

        seq4 = seq1.getComponentByPosition(6)
        pubkey = bit2string(seq4.getComponentByPosition(1))
        return load_der_public_key(pubkey, backend=default_backend())

    def get_prikey(self, passwd):
        pri_filename = os.path.join(self.cert_path, "signPri.key")

        key_file = open(pri_filename, 'rb').read()

        der_seq, remaining = der_decoder.decode(key_file)

        seq1 = der_seq.getComponentByPosition(0)
        self.oid = str(seq1.getComponentByPosition(0))

        seq2 = seq1.getComponentByPosition(1)
        salt = str(seq2.getComponentByPosition(0))
        itercnt = int(seq2.getComponentByPosition(1))

        prikey = str(der_seq.getComponentByPosition(1))

        if self.oid == SEED_CBC_SHA1:
            dk = _pbkdf1(passwd, salt, itercnt)

            key_data = dk[:16]
            s = hasher(dk[16:20])
            iv_data = s.finalize()[:16]

            cipher = Cipher(algorithms.SEED(key_data), modes.CBC(iv_data),
                            backend=default_backend())
            decryptor = cipher.decryptor()
            dec_data = decryptor.update(prikey) + decryptor.finalize()
            padder = padding.PKCS7(128).unpadder()
            dec_data = padder.update(dec_data) + padder.finalize()

            return load_der_private_key(dec_data, password=None,
                                        backend=default_backend())
        else:
            return None

    def sign(self, msg):
        s = self.pri_key.signer(
            asympad.PSS(mgf=asympad.MGF1(hashes.SHA256()),
                        salt_length=asympad.PSS.MAX_LENGTH),
            hashes.SHA256()
        )
        s.update(msg)
        return s.finalize()

    def verify(self, msg, sn):
        v = self.pub_key.verifier(
            sn,
            asympad.PSS(mgf=asympad.MGF1(hashes.SHA256()),
                        salt_length=asympad.PSS.MAX_LENGTH),
            hashes.SHA256()
        )
        v.update(msg)
        try:
            v.verify()
        except InvalidSignature:
            return False
        return True

    def encrypt(self, msg):
        e = self.pub_key.encrypt(
            msg,
            asympad.OAEP(mgf=asympad.MGF1(algorithm=hashes.SHA1()),
                         algorithm=hashes.SHA1(),
                         label=None),
        )
        return e

    def decrypt(self, msg):
        d = self.pri_key.decrypt(
            msg,
            asympad.OAEP(mgf=asympad.MGF1(algorithm=hashes.SHA1()),
                         algorithm=hashes.SHA1(),
                         label=None)
        )
        return d


def select_cert():
    cert_list = []

    def adddir(a, d, f):
        l = os.path.split(d)[1]
        if l[:3] == "cn=":
            cert_list.append(d)

    def _df(d):
        return "20" + d[:2] + "/" + d[2:4] + "/" + d[4:6] + " " + \
               d[6:8] + ":" + d[8:10] + ":" + d[10:12]

    os.path.walk(NPKI_PATH, adddir, None)
    i = 1
    for p in cert_list:
        cert = KoCertificate(p)
        print "%2d: %s (valid: %s ~ %s)" % \
            (i, cert.owner, _df(cert.valid_date[0]), _df(cert.valid_date[1]))
        i += 1

    print
    sel = raw_input("Select Certificate: ")
    return cert_list[int(sel) - 1]


if __name__ == "__main__":
    selpath = select_cert()

    passwd = getpass.getpass("Password: ")
    print

    cert = KoCertificate(selpath, passwd)

    print "OID: ", cert.oid
    print "OWNER: ", cert.owner
    print "VALID: ", cert.valid_date
    print

    s = cert.sign('some test text')
    print "Sign verify:", cert.verify('some test text', s)

    ss = cert.encrypt('some more test text')
    print "Decrypt:", cert.decrypt(ss)