tpm2_smoke.py

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


class SmokeTest(unittest.TestCase):
    def setUp(self):
        self.client = tpm2.Client()
        self.root_key = self.client.create_root_key()

    def tearDown(self):
        self.client.flush_context(self.root_key)
        self.client.close()

    def test_seal_with_auth(self):
        data = 'X' * 64
        auth = 'A' * 15

        blob = self.client.seal(self.root_key, data, auth, None)
        result = self.client.unseal(self.root_key, blob, auth, None)
        self.assertEqual(data, result)

    def test_seal_with_policy(self):
        handle = self.client.start_auth_session(tpm2.TPM2_SE_TRIAL)

        data = 'X' * 64
        auth = 'A' * 15
        pcrs = [16]

        try:
            self.client.policy_pcr(handle, pcrs)
            self.client.policy_password(handle)

            policy_dig = self.client.get_policy_digest(handle)
        finally:
            self.client.flush_context(handle)

        blob = self.client.seal(self.root_key, data, auth, policy_dig)

        handle = self.client.start_auth_session(tpm2.TPM2_SE_POLICY)

        try:
            self.client.policy_pcr(handle, pcrs)
            self.client.policy_password(handle)

            result = self.client.unseal(self.root_key, blob, auth, handle)
        except:
            self.client.flush_context(handle)
            raise

        self.assertEqual(data, result)

    def test_unseal_with_wrong_auth(self):
        data = 'X' * 64
        auth = 'A' * 20
        rc = 0

        blob = self.client.seal(self.root_key, data, auth, None)
        try:
            result = self.client.unseal(self.root_key, blob, auth[:-1] + 'B', None)
        except ProtocolError, e:
            rc = e.rc

        self.assertEqual(rc, tpm2.TPM2_RC_AUTH_FAIL)

    def test_unseal_with_wrong_policy(self):
        handle = self.client.start_auth_session(tpm2.TPM2_SE_TRIAL)

        data = 'X' * 64
        auth = 'A' * 17
        pcrs = [16]

        try:
            self.client.policy_pcr(handle, pcrs)
            self.client.policy_password(handle)

            policy_dig = self.client.get_policy_digest(handle)
        finally:
            self.client.flush_context(handle)

        blob = self.client.seal(self.root_key, data, auth, policy_dig)

        # Extend first a PCR that is not part of the policy and try to unseal.
        # This should succeed.

        ds = tpm2.get_digest_size(tpm2.TPM2_ALG_SHA1)
        self.client.extend_pcr(1, 'X' * ds)

        handle = self.client.start_auth_session(tpm2.TPM2_SE_POLICY)

        try:
            self.client.policy_pcr(handle, pcrs)
            self.client.policy_password(handle)

            result = self.client.unseal(self.root_key, blob, auth, handle)
        except:
            self.client.flush_context(handle)
            raise

        self.assertEqual(data, result)

        # Then, extend a PCR that is part of the policy and try to unseal.
        # This should fail.
        self.client.extend_pcr(16, 'X' * ds)

        handle = self.client.start_auth_session(tpm2.TPM2_SE_POLICY)

        rc = 0

        try:
            self.client.policy_pcr(handle, pcrs)
            self.client.policy_password(handle)

            result = self.client.unseal(self.root_key, blob, auth, handle)
        except ProtocolError, e:
            rc = e.rc
            self.client.flush_context(handle)
        except:
            self.client.flush_context(handle)
            raise

        self.assertEqual(rc, tpm2.TPM2_RC_POLICY_FAIL)

    def test_seal_with_too_long_auth(self):
        ds = tpm2.get_digest_size(tpm2.TPM2_ALG_SHA1)
        data = 'X' * 64
        auth = 'A' * (ds + 1)

        rc = 0
        try:
            blob = self.client.seal(self.root_key, data, auth, None)
        except ProtocolError, e:
            rc = e.rc

        self.assertEqual(rc, tpm2.TPM2_RC_SIZE)