from ctypes import cdll
import ctypes
import os
import tpm2
lib = cdll.LoadLibrary('libtss.so.0')

class TPM2B_NONCE(ctypes.Structure):
    _fields_ = [("s", ctypes.c_uint16),
                ("b", ctypes.c_uint8 * 128)]

class TPM2B_ENCRYPTED_SECRET(ctypes.Structure):
    _fields_ = [("s", ctypes.c_uint16),
                ("b", ctypes.c_uint8 * 256)]

class TPMT_SYM_DEF(ctypes.Structure):
    _fields_ = [("algorithm", ctypes.c_uint16),
                ("keyBits", ctypes.c_uint16),
                ("mode", ctypes.c_uint16)]

class TPMT_RSA_SCHEME(ctypes.Structure):
    _fields_ = [("scheme", ctypes.c_uint16),
                ("details", ctypes.c_uint8 * 4)]

class TPMS_RSA_PARMS(ctypes.Structure):
    _fields_ = [("symmetric", TPMT_SYM_DEF),
                ("scheme", TPMT_RSA_SCHEME),
                ("keyBits", ctypes.c_uint16),
                ("exponent", ctypes.c_uint32)]

class TPMU_PUBLIC_PARMS(ctypes.Union):
    _fields_ = [("buffer", ctypes.c_uint8 *20),
                ("rsaDetail", TPMS_RSA_PARMS)]


class TPMT_PUBLIC(ctypes.Structure):
    _fields_ = [("Type", ctypes.c_uint16),
                ("nameAlg", ctypes.c_uint16),
                ("objectAttributes", ctypes.c_uint32),
                ("authPolicy", TPM2B_NONCE),
                ("parameters", TPMU_PUBLIC_PARMS),
                ("unique", ctypes.c_uint8 * 258)]

class TPM2B_PUBLIC(ctypes.Structure):
    _fields_ = [("size", ctypes.c_uint16),
                ("publicArea", TPMT_PUBLIC)]

class TPMT_TK_CREATION(ctypes.Structure):
    _fields_ = [("tag", ctypes.c_uint16),
                ("hierarchy", ctypes.c_uint32),
                ("digest", TPM2B_NONCE)]

class TPM2B_SENSITIVE_DATA(ctypes.Structure):
    _fields_ = [("size", ctypes.c_uint16),
                ("buffer", ctypes.c_uint8 * 130)]

class TPMS_SENSITIVE_CREATE(ctypes.Structure):
    _fields_ = [("userAuth", TPM2B_NONCE),
                ("data", TPM2B_SENSITIVE_DATA)]

class TPM2B_SENSITIVE_CREATE(ctypes.Structure):
    _fields_ = [("size", ctypes.c_uint16),
                ("sensitive", TPMS_SENSITIVE_CREATE)]

class TPM2B_PRIVATE(ctypes.Structure):
    _fields_ = [("size", ctypes.c_uint16),
                ("buffer", ctypes.c_uint8 * 1166)]

class TPM2B_CONTEXT_DATA(ctypes.Structure):
    _fields_ = [("size", ctypes.c_uint16),
                ("buffer", ctypes.c_uint8 * 2048)]

class TPMS_CONTEXT(ctypes.Structure):
    _fields_ = [("sequence", ctypes.c_uint64),
                ("savedHandle", ctypes.c_uint32),
                ("hierarchy", ctypes.c_uint32),
                ("contextBlob", TPM2B_CONTEXT_DATA)]

class StartAuthSession_In(ctypes.Structure):
    _fields_ = [("tpmKey", ctypes.c_uint32),
                ("bind", ctypes.c_uint32),
                ("nonce", TPM2B_NONCE),
                ("encryptedSalt", TPM2B_ENCRYPTED_SECRET),
                ("sessionType", ctypes.c_uint8),
                ("symmetric", TPMT_SYM_DEF),
                ("authHash", ctypes.c_uint16)]

class StartAuthSession_Out(ctypes.Structure):
    _fields_ = [("sessionHandle", ctypes.c_uint32),
                ("nonceTPM", TPM2B_NONCE)]

class StartAuthSession_Extra(ctypes.Structure):
    _fields_ = [("bindPassword", ctypes.c_char_p),
                ("salt", TPM2B_NONCE)]

class ReadPublic_In(ctypes.Structure):
    _fields_ = [("objectHandle", ctypes.c_uint32)]

class ReadPublic_Out(ctypes.Structure):
    _fields_ = [("outPublic", TPM2B_PUBLIC),
                ("name", TPM2B_NONCE),
                ("qualifiedName", TPM2B_NONCE)]

class Create_In(ctypes.Structure):
    _fields_ = [("parentHandle", ctypes.c_uint32),
                ("inSensitive", TPM2B_SENSITIVE_CREATE),
                ("inPublic", TPM2B_PUBLIC),
                ("outsideInfo", TPM2B_NONCE),
                ("creationPCR", ctypes.c_uint8 * 24)]

class Create_Out(ctypes.Structure):
    _fields_ = [("outPrivate", TPM2B_PRIVATE),
                ("outPublic", TPM2B_PUBLIC),
                ("creationData", ctypes.c_uint8 * 564),
                ("creationHash", TPM2B_NONCE),
                ("creationTicket", TPMT_TK_CREATION)]

class FlushContext_In(ctypes.Structure):
    _fields_ = [("flushHandle", ctypes.c_uint32)]

class Load_In(ctypes.Structure):
    _fields_ = [("parentHandle", ctypes.c_uint32),
                ("inPrivate", TPM2B_PRIVATE),
                ("inPublic", TPM2B_PUBLIC)]

class Load_Out(ctypes.Structure):
    _fields_ = [("objectHandle", ctypes.c_uint32),
                ("name", TPM2B_NONCE)]

class ObjectChangeAuth_In(ctypes.Structure):
    _fields_ = [("objectHandle", ctypes.c_uint32),
                ("parentHandle", ctypes.c_uint32),
                ("newAuth", TPM2B_NONCE)]

class ObjectChangeAuth_Out(ctypes.Structure):
    _fields_ = [("outPrivate", TPM2B_PRIVATE)]

class ContextSave_In(ctypes.Structure):
    _fields_ = [("saveHandle", ctypes.c_uint32)]

class tpm_error(Exception):

    def __init__(self, rc):
        self.raw_rc = rc
        if ((rc & 0x80) == 0x80):         # RC_FMT1
            rc = rc & 0xbf
        self.rc = rc

    def __str__(self):
        msg = ctypes.c_char_p()
        submsg = ctypes.c_char_p()
        num = ctypes.c_char_p()
        lib.TSS_ResponseCode_toString(ctypes.byref(msg),
                                      ctypes.byref(submsg),
                                      ctypes.byref(num), self.raw_rc);
        return "%s%s%s" % (msg.value, submsg.value, num.value)

class Client:

    def __init__(self):
        self.ctx = ctypes.c_void_p(None)
        os.environ["TPM_DEVICE"] = '/dev/tpms1'
        rc = lib.TSS_Create(ctypes.byref(self.ctx))
        # verify SRK is at 81000001
        self.SRK = 0x81000001
        # we need the public area in the context for salted parameter encryption
        self.read_public(self.SRK)

    def close(self):
        lib.TSS_Delete(self.ctx)

    def TSS_Execute(self, out, inp, extra, ordinal, *sessions):
        rc = lib.TSS_Execute(self.ctx, out, inp,
                             extra, ordinal, *sessions)
        if (rc != 0):
            raise tpm_error(rc)

    def start_session(self, Type, salt_key = tpm2.TPM2_RH_NULL):
        inp = StartAuthSession_In()
        out = StartAuthSession_Out()
        extra = StartAuthSession_Extra()

        inp.sessionType = Type
        inp.authHash = tpm2.TPM2_ALG_SHA256
        inp.tpmKey = salt_key
        inp.symmetric.algorithm = tpm2.TPM2_ALG_AES
        inp.symmetric.keyBits = 128
        inp.symmetric.mode = tpm2.TPM2_ALG_CFB

        self.TSS_Execute(ctypes.byref(out), ctypes.byref(inp),
                         ctypes.byref(extra),
                         tpm2.TPM2_CC_START_AUTH_SESSION,
                         tpm2.TPM2_RH_NULL, None, 0)
        return out.sessionHandle

    def read_public(self, handle):
        inp = ReadPublic_In()
        out = ReadPublic_Out()
        inp.objectHandle = handle

        self.TSS_Execute(ctypes.byref(out), ctypes.byref(inp),
                         None, tpm2.TPM2_CC_READ_PUBLIC,
                         tpm2.TPM2_RH_NULL, None, 0)
        return out

    def flush_context(self, handle):
        inp = FlushContext_In()
        inp.flushHandle = handle

        self.TSS_Execute(None, ctypes.byref(inp), None,
                         tpm2.TPM2_CC_FLUSH_CONTEXT,
                         tpm2.TPM2_RH_NULL, None, 0)

    def create_rsa(self, parent, auth, hmacSession = tpm2.TPM2_RS_PW, hmacCont = 0, encSession = tpm2.TPM2_RH_NULL, encCont = 0):
        inp = Create_In()
        out = Create_Out()

        inp.parentHandle = parent
        if (auth != None):
            lenauth = len(auth)
            inp.inSensitive.sensitive.userAuth.b[0:lenauth] = bytearray(auth)
            inp.inSensitive.sensitive.userAuth.s = lenauth
        inp.inPublic.publicArea.Type = tpm2.TPM2_ALG_RSA
        inp.inPublic.publicArea.nameAlg = tpm2.TPM2_ALG_SHA256
        inp.inPublic.publicArea.objectAttributes = tpm2.TPMA_OBJECT_NODA | tpm2.TPMA_OBJECT_DECRYPT | tpm2.TPMA_OBJECT_USERWITHAUTH | tpm2.TPMA_OBJECT_SENSITIVEDATAORIGIN
        inp.inPublic.publicArea.parameters.rsaDetail.symmetric.algorithm = tpm2.TPM2_ALG_NULL
        inp.inPublic.publicArea.parameters.rsaDetail.scheme.scheme = tpm2.TPM2_ALG_NULL
        inp.inPublic.publicArea.parameters.rsaDetail.keyBits = 2048
        inp.inPublic.publicArea.parameters.rsaDetail.exponent = 0
        if (hmacSession == 0):
            hmacSession = tpm2.TPM2_RS_PW

        self.TSS_Execute(ctypes.byref(out), ctypes.byref(inp), None,
                         tpm2.TPM2_CC_CREATE,
                         hmacSession, None, hmacCont,
                         # 0x20 is TPMA_OBJECT_DECRYPT
                         encSession, None, encCont | 0x20,
                         tpm2.TPM2_RH_NULL, None, 0);
        return out

    def load(self, parent, private, public, auth, hmacSession = tpm2.TPM2_RS_PW, hmacCont = 0):
        inp = Load_In()
        out = Load_Out()

        inp.parentHandle = parent
        inp.inPrivate = private
        inp.inPublic = public

        self.TSS_Execute(ctypes.byref(out), ctypes.byref(inp), None,
                         tpm2.TPM2_CC_LOAD,
                         hmacSession, auth, hmacCont,
                         tpm2.TPM2_RH_NULL, None, 0);

        return out.objectHandle

    def change_auth(self, parent, handle, oldauth, newauth, hmacSession = tpm2.TPM2_RS_PW, hmacCont = 0, encSession = tpm2.TPM2_RH_NULL, encCont = 0):
        inp = ObjectChangeAuth_In()
        out = ObjectChangeAuth_Out()

        inp.parentHandle = parent
        inp.objectHandle = handle
        inp.newAuth.b[0:len(newauth)] = bytearray(newauth)
        inp.newAuth.s = len(newauth)

        self.TSS_Execute(ctypes.byref(out), ctypes.byref(inp), None,
                         tpm2.TPM2_CC_OBJECT_CHANGE_AUTH,
                         hmacSession, oldauth, hmacCont,
                         encSession, None, encCont | 0x20,
                         tpm2.TPM2_RH_NULL, None, 0)

        return out.outPrivate

    def context_save(self, handle):
        inp = ContextSave_In()
        out = TPMS_CONTEXT()

        inp.saveHandle = handle

        self.TSS_Execute(ctypes.byref(out), ctypes.byref(inp), None,
                         tpm2.TPM2_CC_CONTEXT_SAVE,
                         tpm2.TPM2_RH_NULL, None, 0)

        return out

    def context_load(self, inp):
        out = ContextSave_In()

        self.TSS_Execute(ctypes.byref(out), ctypes.byref(inp), None,
                         tpm2.TPM2_CC_CONTEXT_LOAD,
                         tpm2.TPM2_RH_NULL, None, 0)

        return out.saveHandle