import platform
from .TpmTypes import *
from .TpmDevice import *


Owner = TPM_HANDLE(TPM_RH.OWNER);
Endorsement = TPM_HANDLE(TPM_RH.ENDORSEMENT);

NullSymDef = TPMT_SYM_DEF(TPM_ALG_ID.NULL, 0, TPM_ALG_ID.NULL)

class Session:
    def __init__(this,
        sessionHandle = None,
        nonceTpm = None,
        sessionAttributes = TPMA_SESSION.continueSession,
        nonceCaller = None
    ):
        this.SessIn = TPMS_AUTH_COMMAND(sessionHandle, nonceCaller, sessionAttributes);
        this.SessOut = TPMS_AUTH_RESPONSE(nonceTpm, sessionAttributes);

    @staticmethod
    def Pw(authValue = None): # Session
        s = Session();
        s.SessIn.sessionHandle = TPM_HANDLE(TPM_RH.RS_PW);
        s.SessIn.nonce = None;
        s.SessIn.sessionAttributes = TPMA_SESSION.continueSession;
        s.SessIn.hmac = authValue;
        s.SessOut.sessionAttributes = TPMA_SESSION.continueSession;
        return s;
# class Session

NullPwSession = Session.Pw()


class TpmBase(object):
    def __init__(self, useSimulator = False, host = '127.0.0.1', port = 2321):
        if useSimulator:
            self.__device = TpmTcpDevice(host, port)
        elif platform.system() == 'Windows':
            self.__device = TpmTbsDevice()
        else:
            self.__device = TpmLinuxDevice()
        self.__lastResponseCode = TPM_RC.SUCCESS
        self.__lastError = None    # TpmError
        self.enableExceptions(True)
        self.__sessions = None
        self.__cmdBuf = None

    def connect(self):
        try:
            self.__device.connect()
        except Exception as e:
            if isinstance(self.__device, TpmLinuxDevice):
                # It is possible that a user mode TRM from tpm2-tools is running
                self.__device = TpmTcpDevice('127.0.0.1', 2323, True)
                self.__device.connect()
            else:
                raise

    def close(self):
        if self.__device:
            self.__device.close()
            self.__device = None

    @property
    def lastResponseCode(self): 
        return self.__lastResponseCode

    @property
    def lastError(self):
        return self.__lastError

    def allowErrors(self):
        """For the next TPM command invocation, errors will not cause an exception to be thrown
           (use _lastCommandSucceeded or _getLastResponseCode() to check for an error)
        Returns:
            This object (to allow modifier chaining)
        """
        self.__errorsAllowed = True
        return self;

    def enableExceptions(self, enable = True):
        """When exceptions are enabled, errors reported by the TPM or occurred in the TSS (e.g. during 
            an attempt to communicate with the TPM) will result in throwing an exception of TpmError type.
            It will still be possible to use _lastCommandSucceeded(), _getLastResponseCode() methods and
            lastError property to check for an error after the exception is intercepted.
            Note that in contrast to allowErrors() this method affects all subsequent commands. 
        """
        self.__exceptionsEnabled = enable
        self.__errorsAllowed = not enable
    
    def withSession(self, sess):
        """Specifies a single session handle to use with the next command
        Args:
            sess Session handle
        Returns:
            This object (to allow modifier chaining)
        """
        self.__sessions = [sess]
        return self

    def withSessions(self, *sessions):
        """Specifies a single session handle to use with the next command
        Args:
            sessions Up to 3 session handles
        Returns:
            This object (to allow modifier chaining)
        """
        print('withSessions: ' + str(NewPython))
        self.__sessions = list(sessions)
        return self

    @staticmethod
    def __isCommMediumError(respCode):
        """Checks whether the response code is generated by the TSS.JS (i.e. is an extension to the TPM 2.0 spec)
        Args:
            code Response code returned by TSS.JS
        Returns:
            true if the response code was generated in the communication channel between the app and the TPM
        """
        return (respCode & 0xFFFF0000) == 0x80280000

    @staticmethod
    def __cleanResponseCode(rawResponse):
        if TpmBase.__isCommMediumError(rawResponse):
            return TPM_RC(rawResponse)

        if rawResponse & TPM_RC.RC_FMT1:
            mask = TPM_RC.RC_FMT1 | 0x3F
        else:
            mask = TPM_RC.RC_WARN | TPM_RC.RC_VER1 | 0x7F
        return TPM_RC(rawResponse & mask)

    def prepareCmdBuf(this,
        cmdCode,                # TPM_CC
        handles,                # TPM_HANDLE[]
        numAuthHandles          # Number
    ):
        cmdBuf = TpmBuffer()

        if numAuthHandles > 0:
            this.__cmdTag = TPM_ST.SESSIONS
        else:
            this.__cmdTag = TPM_ST.NO_SESSIONS

        cmdBuf.toTpm(this.__cmdTag, 2)
        cmdBuf.toTpm(0, 4)  # to be filled in later
        cmdBuf.toTpm(cmdCode, 4)

        if handles and len(handles) > 0:
            for h in handles:
                if not h:
                    cmdBuf.toTpm(TPM_RH.NULL, 4)
                else:
                    h.toTpm(cmdBuf)

        if numAuthHandles > 0:
            if not this.__sessions:
                this.__sessions = [None] * numAuthHandles
            elif len(this.__sessions) < numAuthHandles:
                this.sessions = this.__sessions + [None] * (numAuthHandles - len(this.__sessions))

            for i in range(numAuthHandles):
                if this.__sessions[i] == None:
                    this.__sessions[i] = Session.Pw()

            # We do not know the size of the authorization area yet.
            # Remember the place to marshal it, ...
            authSizePos = cmdBuf.curPos
            # ... and marshal a placeholder 0 value for now.
            cmdBuf.toTpm(0, 4)

            for sess in this.__sessions:
                sess.SessIn.toTpm(cmdBuf)

            authSize = cmdBuf.curPos - authSizePos - 4
            cmdBuf.buffer[authSizePos : authSizePos + 4] = intToTpm(authSize, 4)

        this.__sessions = None
        this.__lastError = None
        # this.__lastResponseCode 
        return cmdBuf
    # prepareCmdBuf()

    def dispatchCommand(this, commandBuffer):
        # Fill in command buffer size in the command header
        cmdBuf = commandBuffer.buffer
        cmdBuf[2 : 6] = intToTpm(len(cmdBuf), 4)
        this.__cmdBuf = commandBuffer;
        rc = TPM_RC.RETRY
        while rc == TPM_RC.RETRY:
            respBuf = this.__device.dispatchCommand(cmdBuf);
            rc = intFromTpm(respBuf, 6, 4)
            #if isinstance(respBuf, bytes):
            #    print('RESP Type: ' + respBuf.__class__.__name__)
        return TpmBuffer(respBuf)
    # __dispatchCommand()

    @staticmethod
    def __generateErrorResponse(rc):
        respBuf = TpmBuffer();
        respBuf.toTpm(TPM_ST.NO_SESSIONS, 2);
        respBuf.toTpm(10, 4);
        respBuf.toTpm(rc, 4);
        return respBuf;

    def __generateError(this, cmdCode, respCode, errMsg, noThrow):
        err = TpmError(respCode, cmdCode, errMsg)
        if this.__exceptionsEnabled and not noThrow:
            raise(err)
        return err

    def getCmdError(this, cmd):
        rc = this.lastResponseCode
        if rc == TPM_RC.SUCCESS:
            return None
        return  TpmError(rc, cmd, "TPM command {" + str(cmd) + "}" + "failed with response code {" + str(rc) + "}")

    # Returns response parameters size
    def processResponse(this, cmdCode, respBuf):
        noThrow = this.__errorsAllowed
        this.__errorsAllowed = not this.__exceptionsEnabled

        if respBuf.length < 10:
            this.__lastError = TpmError(TPM_RC.TSS_RESP_BUF_TOO_SHORT, cmdCode,
                                        'Response buffer is too short: ' + str(len(respBuf)))
            return 0

        if respBuf.curPos != 0:
            raise(Exception('Response buffer reading position is not properly initialized!'))

        tag = respBuf.fromTpm(2)        # TPM_ST
        respSize = respBuf.fromTpm(4)
        rc = respBuf.fromTpm(4)         # TPM_RC

        this.__lastResponseCode = TpmBase.__cleanResponseCode(rc)

        if (rc == TPM_RC.SUCCESS and tag != int(this.__cmdTag)
        or rc != TPM_RC.SUCCESS and tag != int(TPM_ST.NO_SESSIONS)):
            this.__lastError = TpmError(TPM_RC.TSS_RESP_BUF_INVALID_SESSION_TAG, cmdCode,
                                        'Invalid session tag in the response buffer')
            print('TAG Data: ', rc, tag, this.__cmdTag)
            return 0

        if this.__lastResponseCode != TPM_RC.SUCCESS:
            this.__lastError = TpmError(this.__lastResponseCode, cmdCode,
                                        'Command {' + str(TPM_CC(cmdCode)) +'} failed with error {' + str(this.__lastResponseCode) + '}')
            if not noThrow:
                raise(this.lastError)
            return 0

        retHandle = None
        if (cmdCode == TPM_CC.CreatePrimary
        or cmdCode == TPM_CC.Load
        or cmdCode == TPM_CC.HMAC_Start
        or cmdCode == TPM_CC.ContextLoad
        or cmdCode == TPM_CC.LoadExternal
        or cmdCode == TPM_CC.StartAuthSession
        or cmdCode == TPM_CC.HashSequenceStart
        or cmdCode == TPM_CC.CreateLoaded):
            # Response buffer contains a handle returned by the TPM
            retHandle = respBuf.createFromTpm(TPM_HANDLE);
            #assert(retHandle.handle != 0 and retHandle.handle != TPM_RH.UNASSIGNED);

        # If a response session is present, response buffer contains a field
        # specifying the size of response parameters
        respParamsSize = respBuf.length - respBuf.curPos
        if tag == TPM_ST.SESSIONS:
            respParamsSize = respBuf.fromTpm(4)

        if retHandle:
            # A trick to simplify code gen for returned handles handling
            respBuf.curPos = respBuf.curPos - 4
            retHandle.toTpm(respBuf)
            respBuf.curPos = respBuf.curPos - 4
            respParamsSize += 4

        return respParamsSize
    # processResponse()

# class TpmBase
