tpm2_sessions_smoke.py

from argparse import ArgumentParser
from argparse import FileType
import os
import sys
import tpm2
from tpm2 import ProtocolError
import unittest
import logging
import tss2

pwd1 = "wibble"
pwd2 = "newpassword"

class SessionTest(unittest.TestCase):
    def setUp(self):
        self.c = tss2.Client()

    def tearDown(self):
        self.c.close()

    # open handles until failure.  Return the ones we got
    def open_handles(self):
        ha = []
        try:
            for i in range(0, 10):
                h = self.c.start_session(tpm2.TPM2_SE_HMAC)
                print "Handle is %08x" % h
                ha.append(h)
        except tss2.tpm_error, e:
            if (e.rc != tpm2.TPM2_RC_SESSION_MEMORY):
                raise e
        return ha

    def test_handle_flush_on_space_close(self):
        i = self.open_handles()
        print "Ran out of handles at %d" %len(i)
        self.c.close()
        self.c = tss2.Client()
        # closing and reopening a space session should clear out our handles
        j = self.open_handles()
        print "Ran out of handles at %d" %len(j)
        self.assertNotEqual(len(i), 0)
        self.assertEqual(len(i), len(j))

    def test_flush(self):
        i = self.open_handles()
        print "opened %d handles" % len(i)
        self.c.flush_context(i[0])
        self.c.flush_context(i[1])
        i = self.open_handles()
        self.assertEqual(len(i), 2);

    def test_session_consumption(self):
        self.c.read_public(self.c.SRK)
        # authorization hmac session
        hmac = self.c.start_session(tpm2.TPM2_SE_HMAC)
        # parameter encryption session
        enc = self.c.start_session(tpm2.TPM2_SE_HMAC, self.c.SRK)
        # fill all remaing handles
        i = self.open_handles()
        # create rsa key continuing both hmac and encryption sessions
        self.c.create_rsa(self.c.SRK, pwd1, hmac, 1, enc, 1)
        # should be no handles left
        i = self.open_handles()
        self.assertEqual(len(i),0)
        # now create rsa key continuing hmac and consuming encryption
        k = self.c.create_rsa(self.c.SRK, pwd1, hmac, 1, enc, 0)
        # now should be one handle remaining
        i = self.open_handles()
        self.assertEqual(len(i),1)
        self.c.flush_context(i[0])
        # check the hmac continuation actually works
        k = self.c.load(self.c.SRK, k.outPrivate, k.outPublic, None)
        print "Loaded key at handle %x" %k
        # and finally verify with an authenticated encrypted operation
        # consuming both handles
        enc = self.c.start_session(tpm2.TPM2_SE_HMAC, k)
        self.c.change_auth(self.c.SRK, k, pwd1, pwd2, hmac, 0, enc, 0)
        i = self.open_handles()
        self.assertEqual(len(i), 2)
        
        


if __name__ == '__main__':
    unittest.main()