162306a36Sopenharmony_ci# SPDX-License-Identifier: (GPL-2.0 OR BSD-3-Clause)
262306a36Sopenharmony_ci
362306a36Sopenharmony_ciimport hashlib
462306a36Sopenharmony_ciimport os
562306a36Sopenharmony_ciimport socket
662306a36Sopenharmony_ciimport struct
762306a36Sopenharmony_ciimport sys
862306a36Sopenharmony_ciimport unittest
962306a36Sopenharmony_ciimport fcntl
1062306a36Sopenharmony_ciimport select
1162306a36Sopenharmony_ci
1262306a36Sopenharmony_ciTPM2_ST_NO_SESSIONS = 0x8001
1362306a36Sopenharmony_ciTPM2_ST_SESSIONS = 0x8002
1462306a36Sopenharmony_ci
1562306a36Sopenharmony_ciTPM2_CC_FIRST = 0x01FF
1662306a36Sopenharmony_ci
1762306a36Sopenharmony_ciTPM2_CC_CREATE_PRIMARY = 0x0131
1862306a36Sopenharmony_ciTPM2_CC_DICTIONARY_ATTACK_LOCK_RESET = 0x0139
1962306a36Sopenharmony_ciTPM2_CC_CREATE = 0x0153
2062306a36Sopenharmony_ciTPM2_CC_LOAD = 0x0157
2162306a36Sopenharmony_ciTPM2_CC_UNSEAL = 0x015E
2262306a36Sopenharmony_ciTPM2_CC_FLUSH_CONTEXT = 0x0165
2362306a36Sopenharmony_ciTPM2_CC_START_AUTH_SESSION = 0x0176
2462306a36Sopenharmony_ciTPM2_CC_GET_CAPABILITY	= 0x017A
2562306a36Sopenharmony_ciTPM2_CC_GET_RANDOM = 0x017B
2662306a36Sopenharmony_ciTPM2_CC_PCR_READ = 0x017E
2762306a36Sopenharmony_ciTPM2_CC_POLICY_PCR = 0x017F
2862306a36Sopenharmony_ciTPM2_CC_PCR_EXTEND = 0x0182
2962306a36Sopenharmony_ciTPM2_CC_POLICY_PASSWORD = 0x018C
3062306a36Sopenharmony_ciTPM2_CC_POLICY_GET_DIGEST = 0x0189
3162306a36Sopenharmony_ci
3262306a36Sopenharmony_ciTPM2_SE_POLICY = 0x01
3362306a36Sopenharmony_ciTPM2_SE_TRIAL = 0x03
3462306a36Sopenharmony_ci
3562306a36Sopenharmony_ciTPM2_ALG_RSA = 0x0001
3662306a36Sopenharmony_ciTPM2_ALG_SHA1 = 0x0004
3762306a36Sopenharmony_ciTPM2_ALG_AES = 0x0006
3862306a36Sopenharmony_ciTPM2_ALG_KEYEDHASH = 0x0008
3962306a36Sopenharmony_ciTPM2_ALG_SHA256 = 0x000B
4062306a36Sopenharmony_ciTPM2_ALG_NULL = 0x0010
4162306a36Sopenharmony_ciTPM2_ALG_CBC = 0x0042
4262306a36Sopenharmony_ciTPM2_ALG_CFB = 0x0043
4362306a36Sopenharmony_ci
4462306a36Sopenharmony_ciTPM2_RH_OWNER = 0x40000001
4562306a36Sopenharmony_ciTPM2_RH_NULL = 0x40000007
4662306a36Sopenharmony_ciTPM2_RH_LOCKOUT = 0x4000000A
4762306a36Sopenharmony_ciTPM2_RS_PW = 0x40000009
4862306a36Sopenharmony_ci
4962306a36Sopenharmony_ciTPM2_RC_SIZE            = 0x01D5
5062306a36Sopenharmony_ciTPM2_RC_AUTH_FAIL       = 0x098E
5162306a36Sopenharmony_ciTPM2_RC_POLICY_FAIL     = 0x099D
5262306a36Sopenharmony_ciTPM2_RC_COMMAND_CODE    = 0x0143
5362306a36Sopenharmony_ci
5462306a36Sopenharmony_ciTSS2_RC_LAYER_SHIFT = 16
5562306a36Sopenharmony_ciTSS2_RESMGR_TPM_RC_LAYER = (11 << TSS2_RC_LAYER_SHIFT)
5662306a36Sopenharmony_ci
5762306a36Sopenharmony_ciTPM2_CAP_HANDLES = 0x00000001
5862306a36Sopenharmony_ciTPM2_CAP_COMMANDS = 0x00000002
5962306a36Sopenharmony_ciTPM2_CAP_PCRS = 0x00000005
6062306a36Sopenharmony_ciTPM2_CAP_TPM_PROPERTIES = 0x00000006
6162306a36Sopenharmony_ci
6262306a36Sopenharmony_ciTPM2_PT_FIXED = 0x100
6362306a36Sopenharmony_ciTPM2_PT_TOTAL_COMMANDS = TPM2_PT_FIXED + 41
6462306a36Sopenharmony_ci
6562306a36Sopenharmony_ciHR_SHIFT = 24
6662306a36Sopenharmony_ciHR_LOADED_SESSION = 0x02000000
6762306a36Sopenharmony_ciHR_TRANSIENT = 0x80000000
6862306a36Sopenharmony_ci
6962306a36Sopenharmony_ciSHA1_DIGEST_SIZE = 20
7062306a36Sopenharmony_ciSHA256_DIGEST_SIZE = 32
7162306a36Sopenharmony_ci
7262306a36Sopenharmony_ciTPM2_VER0_ERRORS = {
7362306a36Sopenharmony_ci    0x000: "TPM_RC_SUCCESS",
7462306a36Sopenharmony_ci    0x030: "TPM_RC_BAD_TAG",
7562306a36Sopenharmony_ci}
7662306a36Sopenharmony_ci
7762306a36Sopenharmony_ciTPM2_VER1_ERRORS = {
7862306a36Sopenharmony_ci    0x000: "TPM_RC_FAILURE",
7962306a36Sopenharmony_ci    0x001: "TPM_RC_FAILURE",
8062306a36Sopenharmony_ci    0x003: "TPM_RC_SEQUENCE",
8162306a36Sopenharmony_ci    0x00B: "TPM_RC_PRIVATE",
8262306a36Sopenharmony_ci    0x019: "TPM_RC_HMAC",
8362306a36Sopenharmony_ci    0x020: "TPM_RC_DISABLED",
8462306a36Sopenharmony_ci    0x021: "TPM_RC_EXCLUSIVE",
8562306a36Sopenharmony_ci    0x024: "TPM_RC_AUTH_TYPE",
8662306a36Sopenharmony_ci    0x025: "TPM_RC_AUTH_MISSING",
8762306a36Sopenharmony_ci    0x026: "TPM_RC_POLICY",
8862306a36Sopenharmony_ci    0x027: "TPM_RC_PCR",
8962306a36Sopenharmony_ci    0x028: "TPM_RC_PCR_CHANGED",
9062306a36Sopenharmony_ci    0x02D: "TPM_RC_UPGRADE",
9162306a36Sopenharmony_ci    0x02E: "TPM_RC_TOO_MANY_CONTEXTS",
9262306a36Sopenharmony_ci    0x02F: "TPM_RC_AUTH_UNAVAILABLE",
9362306a36Sopenharmony_ci    0x030: "TPM_RC_REBOOT",
9462306a36Sopenharmony_ci    0x031: "TPM_RC_UNBALANCED",
9562306a36Sopenharmony_ci    0x042: "TPM_RC_COMMAND_SIZE",
9662306a36Sopenharmony_ci    0x043: "TPM_RC_COMMAND_CODE",
9762306a36Sopenharmony_ci    0x044: "TPM_RC_AUTHSIZE",
9862306a36Sopenharmony_ci    0x045: "TPM_RC_AUTH_CONTEXT",
9962306a36Sopenharmony_ci    0x046: "TPM_RC_NV_RANGE",
10062306a36Sopenharmony_ci    0x047: "TPM_RC_NV_SIZE",
10162306a36Sopenharmony_ci    0x048: "TPM_RC_NV_LOCKED",
10262306a36Sopenharmony_ci    0x049: "TPM_RC_NV_AUTHORIZATION",
10362306a36Sopenharmony_ci    0x04A: "TPM_RC_NV_UNINITIALIZED",
10462306a36Sopenharmony_ci    0x04B: "TPM_RC_NV_SPACE",
10562306a36Sopenharmony_ci    0x04C: "TPM_RC_NV_DEFINED",
10662306a36Sopenharmony_ci    0x050: "TPM_RC_BAD_CONTEXT",
10762306a36Sopenharmony_ci    0x051: "TPM_RC_CPHASH",
10862306a36Sopenharmony_ci    0x052: "TPM_RC_PARENT",
10962306a36Sopenharmony_ci    0x053: "TPM_RC_NEEDS_TEST",
11062306a36Sopenharmony_ci    0x054: "TPM_RC_NO_RESULT",
11162306a36Sopenharmony_ci    0x055: "TPM_RC_SENSITIVE",
11262306a36Sopenharmony_ci    0x07F: "RC_MAX_FM0",
11362306a36Sopenharmony_ci}
11462306a36Sopenharmony_ci
11562306a36Sopenharmony_ciTPM2_FMT1_ERRORS = {
11662306a36Sopenharmony_ci    0x001: "TPM_RC_ASYMMETRIC",
11762306a36Sopenharmony_ci    0x002: "TPM_RC_ATTRIBUTES",
11862306a36Sopenharmony_ci    0x003: "TPM_RC_HASH",
11962306a36Sopenharmony_ci    0x004: "TPM_RC_VALUE",
12062306a36Sopenharmony_ci    0x005: "TPM_RC_HIERARCHY",
12162306a36Sopenharmony_ci    0x007: "TPM_RC_KEY_SIZE",
12262306a36Sopenharmony_ci    0x008: "TPM_RC_MGF",
12362306a36Sopenharmony_ci    0x009: "TPM_RC_MODE",
12462306a36Sopenharmony_ci    0x00A: "TPM_RC_TYPE",
12562306a36Sopenharmony_ci    0x00B: "TPM_RC_HANDLE",
12662306a36Sopenharmony_ci    0x00C: "TPM_RC_KDF",
12762306a36Sopenharmony_ci    0x00D: "TPM_RC_RANGE",
12862306a36Sopenharmony_ci    0x00E: "TPM_RC_AUTH_FAIL",
12962306a36Sopenharmony_ci    0x00F: "TPM_RC_NONCE",
13062306a36Sopenharmony_ci    0x010: "TPM_RC_PP",
13162306a36Sopenharmony_ci    0x012: "TPM_RC_SCHEME",
13262306a36Sopenharmony_ci    0x015: "TPM_RC_SIZE",
13362306a36Sopenharmony_ci    0x016: "TPM_RC_SYMMETRIC",
13462306a36Sopenharmony_ci    0x017: "TPM_RC_TAG",
13562306a36Sopenharmony_ci    0x018: "TPM_RC_SELECTOR",
13662306a36Sopenharmony_ci    0x01A: "TPM_RC_INSUFFICIENT",
13762306a36Sopenharmony_ci    0x01B: "TPM_RC_SIGNATURE",
13862306a36Sopenharmony_ci    0x01C: "TPM_RC_KEY",
13962306a36Sopenharmony_ci    0x01D: "TPM_RC_POLICY_FAIL",
14062306a36Sopenharmony_ci    0x01F: "TPM_RC_INTEGRITY",
14162306a36Sopenharmony_ci    0x020: "TPM_RC_TICKET",
14262306a36Sopenharmony_ci    0x021: "TPM_RC_RESERVED_BITS",
14362306a36Sopenharmony_ci    0x022: "TPM_RC_BAD_AUTH",
14462306a36Sopenharmony_ci    0x023: "TPM_RC_EXPIRED",
14562306a36Sopenharmony_ci    0x024: "TPM_RC_POLICY_CC",
14662306a36Sopenharmony_ci    0x025: "TPM_RC_BINDING",
14762306a36Sopenharmony_ci    0x026: "TPM_RC_CURVE",
14862306a36Sopenharmony_ci    0x027: "TPM_RC_ECC_POINT",
14962306a36Sopenharmony_ci}
15062306a36Sopenharmony_ci
15162306a36Sopenharmony_ciTPM2_WARN_ERRORS = {
15262306a36Sopenharmony_ci    0x001: "TPM_RC_CONTEXT_GAP",
15362306a36Sopenharmony_ci    0x002: "TPM_RC_OBJECT_MEMORY",
15462306a36Sopenharmony_ci    0x003: "TPM_RC_SESSION_MEMORY",
15562306a36Sopenharmony_ci    0x004: "TPM_RC_MEMORY",
15662306a36Sopenharmony_ci    0x005: "TPM_RC_SESSION_HANDLES",
15762306a36Sopenharmony_ci    0x006: "TPM_RC_OBJECT_HANDLES",
15862306a36Sopenharmony_ci    0x007: "TPM_RC_LOCALITY",
15962306a36Sopenharmony_ci    0x008: "TPM_RC_YIELDED",
16062306a36Sopenharmony_ci    0x009: "TPM_RC_CANCELED",
16162306a36Sopenharmony_ci    0x00A: "TPM_RC_TESTING",
16262306a36Sopenharmony_ci    0x010: "TPM_RC_REFERENCE_H0",
16362306a36Sopenharmony_ci    0x011: "TPM_RC_REFERENCE_H1",
16462306a36Sopenharmony_ci    0x012: "TPM_RC_REFERENCE_H2",
16562306a36Sopenharmony_ci    0x013: "TPM_RC_REFERENCE_H3",
16662306a36Sopenharmony_ci    0x014: "TPM_RC_REFERENCE_H4",
16762306a36Sopenharmony_ci    0x015: "TPM_RC_REFERENCE_H5",
16862306a36Sopenharmony_ci    0x016: "TPM_RC_REFERENCE_H6",
16962306a36Sopenharmony_ci    0x018: "TPM_RC_REFERENCE_S0",
17062306a36Sopenharmony_ci    0x019: "TPM_RC_REFERENCE_S1",
17162306a36Sopenharmony_ci    0x01A: "TPM_RC_REFERENCE_S2",
17262306a36Sopenharmony_ci    0x01B: "TPM_RC_REFERENCE_S3",
17362306a36Sopenharmony_ci    0x01C: "TPM_RC_REFERENCE_S4",
17462306a36Sopenharmony_ci    0x01D: "TPM_RC_REFERENCE_S5",
17562306a36Sopenharmony_ci    0x01E: "TPM_RC_REFERENCE_S6",
17662306a36Sopenharmony_ci    0x020: "TPM_RC_NV_RATE",
17762306a36Sopenharmony_ci    0x021: "TPM_RC_LOCKOUT",
17862306a36Sopenharmony_ci    0x022: "TPM_RC_RETRY",
17962306a36Sopenharmony_ci    0x023: "TPM_RC_NV_UNAVAILABLE",
18062306a36Sopenharmony_ci    0x7F: "TPM_RC_NOT_USED",
18162306a36Sopenharmony_ci}
18262306a36Sopenharmony_ci
18362306a36Sopenharmony_ciRC_VER1 = 0x100
18462306a36Sopenharmony_ciRC_FMT1 = 0x080
18562306a36Sopenharmony_ciRC_WARN = 0x900
18662306a36Sopenharmony_ci
18762306a36Sopenharmony_ciALG_DIGEST_SIZE_MAP = {
18862306a36Sopenharmony_ci    TPM2_ALG_SHA1: SHA1_DIGEST_SIZE,
18962306a36Sopenharmony_ci    TPM2_ALG_SHA256: SHA256_DIGEST_SIZE,
19062306a36Sopenharmony_ci}
19162306a36Sopenharmony_ci
19262306a36Sopenharmony_ciALG_HASH_FUNCTION_MAP = {
19362306a36Sopenharmony_ci    TPM2_ALG_SHA1: hashlib.sha1,
19462306a36Sopenharmony_ci    TPM2_ALG_SHA256: hashlib.sha256
19562306a36Sopenharmony_ci}
19662306a36Sopenharmony_ci
19762306a36Sopenharmony_ciNAME_ALG_MAP = {
19862306a36Sopenharmony_ci    "sha1": TPM2_ALG_SHA1,
19962306a36Sopenharmony_ci    "sha256": TPM2_ALG_SHA256,
20062306a36Sopenharmony_ci}
20162306a36Sopenharmony_ci
20262306a36Sopenharmony_ci
20362306a36Sopenharmony_ciclass UnknownAlgorithmIdError(Exception):
20462306a36Sopenharmony_ci    def __init__(self, alg):
20562306a36Sopenharmony_ci        self.alg = alg
20662306a36Sopenharmony_ci
20762306a36Sopenharmony_ci    def __str__(self):
20862306a36Sopenharmony_ci        return '0x%0x' % (alg)
20962306a36Sopenharmony_ci
21062306a36Sopenharmony_ci
21162306a36Sopenharmony_ciclass UnknownAlgorithmNameError(Exception):
21262306a36Sopenharmony_ci    def __init__(self, name):
21362306a36Sopenharmony_ci        self.name = name
21462306a36Sopenharmony_ci
21562306a36Sopenharmony_ci    def __str__(self):
21662306a36Sopenharmony_ci        return name
21762306a36Sopenharmony_ci
21862306a36Sopenharmony_ci
21962306a36Sopenharmony_ciclass UnknownPCRBankError(Exception):
22062306a36Sopenharmony_ci    def __init__(self, alg):
22162306a36Sopenharmony_ci        self.alg = alg
22262306a36Sopenharmony_ci
22362306a36Sopenharmony_ci    def __str__(self):
22462306a36Sopenharmony_ci        return '0x%0x' % (alg)
22562306a36Sopenharmony_ci
22662306a36Sopenharmony_ci
22762306a36Sopenharmony_ciclass ProtocolError(Exception):
22862306a36Sopenharmony_ci    def __init__(self, cc, rc):
22962306a36Sopenharmony_ci        self.cc = cc
23062306a36Sopenharmony_ci        self.rc = rc
23162306a36Sopenharmony_ci
23262306a36Sopenharmony_ci        if (rc & RC_FMT1) == RC_FMT1:
23362306a36Sopenharmony_ci            self.name = TPM2_FMT1_ERRORS.get(rc & 0x3f, "TPM_RC_UNKNOWN")
23462306a36Sopenharmony_ci        elif (rc & RC_WARN) == RC_WARN:
23562306a36Sopenharmony_ci            self.name = TPM2_WARN_ERRORS.get(rc & 0x7f, "TPM_RC_UNKNOWN")
23662306a36Sopenharmony_ci        elif (rc & RC_VER1) == RC_VER1:
23762306a36Sopenharmony_ci            self.name = TPM2_VER1_ERRORS.get(rc & 0x7f, "TPM_RC_UNKNOWN")
23862306a36Sopenharmony_ci        else:
23962306a36Sopenharmony_ci            self.name = TPM2_VER0_ERRORS.get(rc & 0x7f, "TPM_RC_UNKNOWN")
24062306a36Sopenharmony_ci
24162306a36Sopenharmony_ci    def __str__(self):
24262306a36Sopenharmony_ci        if self.cc:
24362306a36Sopenharmony_ci            return '%s: cc=0x%08x, rc=0x%08x' % (self.name, self.cc, self.rc)
24462306a36Sopenharmony_ci        else:
24562306a36Sopenharmony_ci            return '%s: rc=0x%08x' % (self.name, self.rc)
24662306a36Sopenharmony_ci
24762306a36Sopenharmony_ci
24862306a36Sopenharmony_ciclass AuthCommand(object):
24962306a36Sopenharmony_ci    """TPMS_AUTH_COMMAND"""
25062306a36Sopenharmony_ci
25162306a36Sopenharmony_ci    def __init__(self, session_handle=TPM2_RS_PW, nonce=bytes(),
25262306a36Sopenharmony_ci                 session_attributes=0, hmac=bytes()):
25362306a36Sopenharmony_ci        self.session_handle = session_handle
25462306a36Sopenharmony_ci        self.nonce = nonce
25562306a36Sopenharmony_ci        self.session_attributes = session_attributes
25662306a36Sopenharmony_ci        self.hmac = hmac
25762306a36Sopenharmony_ci
25862306a36Sopenharmony_ci    def __bytes__(self):
25962306a36Sopenharmony_ci        fmt = '>I H%us B H%us' % (len(self.nonce), len(self.hmac))
26062306a36Sopenharmony_ci        return struct.pack(fmt, self.session_handle, len(self.nonce),
26162306a36Sopenharmony_ci                           self.nonce, self.session_attributes, len(self.hmac),
26262306a36Sopenharmony_ci                           self.hmac)
26362306a36Sopenharmony_ci
26462306a36Sopenharmony_ci    def __len__(self):
26562306a36Sopenharmony_ci        fmt = '>I H%us B H%us' % (len(self.nonce), len(self.hmac))
26662306a36Sopenharmony_ci        return struct.calcsize(fmt)
26762306a36Sopenharmony_ci
26862306a36Sopenharmony_ci
26962306a36Sopenharmony_ciclass SensitiveCreate(object):
27062306a36Sopenharmony_ci    """TPMS_SENSITIVE_CREATE"""
27162306a36Sopenharmony_ci
27262306a36Sopenharmony_ci    def __init__(self, user_auth=bytes(), data=bytes()):
27362306a36Sopenharmony_ci        self.user_auth = user_auth
27462306a36Sopenharmony_ci        self.data = data
27562306a36Sopenharmony_ci
27662306a36Sopenharmony_ci    def __bytes__(self):
27762306a36Sopenharmony_ci        fmt = '>H%us H%us' % (len(self.user_auth), len(self.data))
27862306a36Sopenharmony_ci        return struct.pack(fmt, len(self.user_auth), self.user_auth,
27962306a36Sopenharmony_ci                           len(self.data), self.data)
28062306a36Sopenharmony_ci
28162306a36Sopenharmony_ci    def __len__(self):
28262306a36Sopenharmony_ci        fmt = '>H%us H%us' % (len(self.user_auth), len(self.data))
28362306a36Sopenharmony_ci        return struct.calcsize(fmt)
28462306a36Sopenharmony_ci
28562306a36Sopenharmony_ci
28662306a36Sopenharmony_ciclass Public(object):
28762306a36Sopenharmony_ci    """TPMT_PUBLIC"""
28862306a36Sopenharmony_ci
28962306a36Sopenharmony_ci    FIXED_TPM = (1 << 1)
29062306a36Sopenharmony_ci    FIXED_PARENT = (1 << 4)
29162306a36Sopenharmony_ci    SENSITIVE_DATA_ORIGIN = (1 << 5)
29262306a36Sopenharmony_ci    USER_WITH_AUTH = (1 << 6)
29362306a36Sopenharmony_ci    RESTRICTED = (1 << 16)
29462306a36Sopenharmony_ci    DECRYPT = (1 << 17)
29562306a36Sopenharmony_ci
29662306a36Sopenharmony_ci    def __fmt(self):
29762306a36Sopenharmony_ci        return '>HHIH%us%usH%us' % \
29862306a36Sopenharmony_ci            (len(self.auth_policy), len(self.parameters), len(self.unique))
29962306a36Sopenharmony_ci
30062306a36Sopenharmony_ci    def __init__(self, object_type, name_alg, object_attributes,
30162306a36Sopenharmony_ci                 auth_policy=bytes(), parameters=bytes(),
30262306a36Sopenharmony_ci                 unique=bytes()):
30362306a36Sopenharmony_ci        self.object_type = object_type
30462306a36Sopenharmony_ci        self.name_alg = name_alg
30562306a36Sopenharmony_ci        self.object_attributes = object_attributes
30662306a36Sopenharmony_ci        self.auth_policy = auth_policy
30762306a36Sopenharmony_ci        self.parameters = parameters
30862306a36Sopenharmony_ci        self.unique = unique
30962306a36Sopenharmony_ci
31062306a36Sopenharmony_ci    def __bytes__(self):
31162306a36Sopenharmony_ci        return struct.pack(self.__fmt(),
31262306a36Sopenharmony_ci                           self.object_type,
31362306a36Sopenharmony_ci                           self.name_alg,
31462306a36Sopenharmony_ci                           self.object_attributes,
31562306a36Sopenharmony_ci                           len(self.auth_policy),
31662306a36Sopenharmony_ci                           self.auth_policy,
31762306a36Sopenharmony_ci                           self.parameters,
31862306a36Sopenharmony_ci                           len(self.unique),
31962306a36Sopenharmony_ci                           self.unique)
32062306a36Sopenharmony_ci
32162306a36Sopenharmony_ci    def __len__(self):
32262306a36Sopenharmony_ci        return struct.calcsize(self.__fmt())
32362306a36Sopenharmony_ci
32462306a36Sopenharmony_ci
32562306a36Sopenharmony_cidef get_digest_size(alg):
32662306a36Sopenharmony_ci    ds = ALG_DIGEST_SIZE_MAP.get(alg)
32762306a36Sopenharmony_ci    if not ds:
32862306a36Sopenharmony_ci        raise UnknownAlgorithmIdError(alg)
32962306a36Sopenharmony_ci    return ds
33062306a36Sopenharmony_ci
33162306a36Sopenharmony_ci
33262306a36Sopenharmony_cidef get_hash_function(alg):
33362306a36Sopenharmony_ci    f = ALG_HASH_FUNCTION_MAP.get(alg)
33462306a36Sopenharmony_ci    if not f:
33562306a36Sopenharmony_ci        raise UnknownAlgorithmIdError(alg)
33662306a36Sopenharmony_ci    return f
33762306a36Sopenharmony_ci
33862306a36Sopenharmony_ci
33962306a36Sopenharmony_cidef get_algorithm(name):
34062306a36Sopenharmony_ci    alg = NAME_ALG_MAP.get(name)
34162306a36Sopenharmony_ci    if not alg:
34262306a36Sopenharmony_ci        raise UnknownAlgorithmNameError(name)
34362306a36Sopenharmony_ci    return alg
34462306a36Sopenharmony_ci
34562306a36Sopenharmony_ci
34662306a36Sopenharmony_cidef hex_dump(d):
34762306a36Sopenharmony_ci    d = [format(x, '02x') for x in d]
34862306a36Sopenharmony_ci    d = [d[i: i + 16] for i in range(0, len(d), 16)]
34962306a36Sopenharmony_ci    d = [' '.join(x) for x in d]
35062306a36Sopenharmony_ci    d = os.linesep.join(d)
35162306a36Sopenharmony_ci
35262306a36Sopenharmony_ci    return d
35362306a36Sopenharmony_ci
35462306a36Sopenharmony_ciclass Client:
35562306a36Sopenharmony_ci    FLAG_DEBUG = 0x01
35662306a36Sopenharmony_ci    FLAG_SPACE = 0x02
35762306a36Sopenharmony_ci    FLAG_NONBLOCK = 0x04
35862306a36Sopenharmony_ci    TPM_IOC_NEW_SPACE = 0xa200
35962306a36Sopenharmony_ci
36062306a36Sopenharmony_ci    def __init__(self, flags = 0):
36162306a36Sopenharmony_ci        self.flags = flags
36262306a36Sopenharmony_ci
36362306a36Sopenharmony_ci        if (self.flags & Client.FLAG_SPACE) == 0:
36462306a36Sopenharmony_ci            self.tpm = open('/dev/tpm0', 'r+b', buffering=0)
36562306a36Sopenharmony_ci        else:
36662306a36Sopenharmony_ci            self.tpm = open('/dev/tpmrm0', 'r+b', buffering=0)
36762306a36Sopenharmony_ci
36862306a36Sopenharmony_ci        if (self.flags & Client.FLAG_NONBLOCK):
36962306a36Sopenharmony_ci            flags = fcntl.fcntl(self.tpm, fcntl.F_GETFL)
37062306a36Sopenharmony_ci            flags |= os.O_NONBLOCK
37162306a36Sopenharmony_ci            fcntl.fcntl(self.tpm, fcntl.F_SETFL, flags)
37262306a36Sopenharmony_ci            self.tpm_poll = select.poll()
37362306a36Sopenharmony_ci
37462306a36Sopenharmony_ci    def __del__(self):
37562306a36Sopenharmony_ci        if self.tpm:
37662306a36Sopenharmony_ci            self.tpm.close()
37762306a36Sopenharmony_ci
37862306a36Sopenharmony_ci    def close(self):
37962306a36Sopenharmony_ci        self.tpm.close()
38062306a36Sopenharmony_ci
38162306a36Sopenharmony_ci    def send_cmd(self, cmd):
38262306a36Sopenharmony_ci        self.tpm.write(cmd)
38362306a36Sopenharmony_ci
38462306a36Sopenharmony_ci        if (self.flags & Client.FLAG_NONBLOCK):
38562306a36Sopenharmony_ci            self.tpm_poll.register(self.tpm, select.POLLIN)
38662306a36Sopenharmony_ci            self.tpm_poll.poll(10000)
38762306a36Sopenharmony_ci
38862306a36Sopenharmony_ci        rsp = self.tpm.read()
38962306a36Sopenharmony_ci
39062306a36Sopenharmony_ci        if (self.flags & Client.FLAG_NONBLOCK):
39162306a36Sopenharmony_ci            self.tpm_poll.unregister(self.tpm)
39262306a36Sopenharmony_ci
39362306a36Sopenharmony_ci        if (self.flags & Client.FLAG_DEBUG) != 0:
39462306a36Sopenharmony_ci            sys.stderr.write('cmd' + os.linesep)
39562306a36Sopenharmony_ci            sys.stderr.write(hex_dump(cmd) + os.linesep)
39662306a36Sopenharmony_ci            sys.stderr.write('rsp' + os.linesep)
39762306a36Sopenharmony_ci            sys.stderr.write(hex_dump(rsp) + os.linesep)
39862306a36Sopenharmony_ci
39962306a36Sopenharmony_ci        rc = struct.unpack('>I', rsp[6:10])[0]
40062306a36Sopenharmony_ci        if rc != 0:
40162306a36Sopenharmony_ci            cc = struct.unpack('>I', cmd[6:10])[0]
40262306a36Sopenharmony_ci            raise ProtocolError(cc, rc)
40362306a36Sopenharmony_ci
40462306a36Sopenharmony_ci        return rsp
40562306a36Sopenharmony_ci
40662306a36Sopenharmony_ci    def read_pcr(self, i, bank_alg = TPM2_ALG_SHA1):
40762306a36Sopenharmony_ci        pcrsel_len = max((i >> 3) + 1, 3)
40862306a36Sopenharmony_ci        pcrsel = [0] * pcrsel_len
40962306a36Sopenharmony_ci        pcrsel[i >> 3] = 1 << (i & 7)
41062306a36Sopenharmony_ci        pcrsel = ''.join(map(chr, pcrsel)).encode()
41162306a36Sopenharmony_ci
41262306a36Sopenharmony_ci        fmt = '>HII IHB%us' % (pcrsel_len)
41362306a36Sopenharmony_ci        cmd = struct.pack(fmt,
41462306a36Sopenharmony_ci                          TPM2_ST_NO_SESSIONS,
41562306a36Sopenharmony_ci                          struct.calcsize(fmt),
41662306a36Sopenharmony_ci                          TPM2_CC_PCR_READ,
41762306a36Sopenharmony_ci                          1,
41862306a36Sopenharmony_ci                          bank_alg,
41962306a36Sopenharmony_ci                          pcrsel_len, pcrsel)
42062306a36Sopenharmony_ci
42162306a36Sopenharmony_ci        rsp = self.send_cmd(cmd)
42262306a36Sopenharmony_ci
42362306a36Sopenharmony_ci        pcr_update_cnt, pcr_select_cnt = struct.unpack('>II', rsp[10:18])
42462306a36Sopenharmony_ci        assert pcr_select_cnt == 1
42562306a36Sopenharmony_ci        rsp = rsp[18:]
42662306a36Sopenharmony_ci
42762306a36Sopenharmony_ci        alg2, pcrsel_len2 = struct.unpack('>HB', rsp[:3])
42862306a36Sopenharmony_ci        assert bank_alg == alg2 and pcrsel_len == pcrsel_len2
42962306a36Sopenharmony_ci        rsp = rsp[3 + pcrsel_len:]
43062306a36Sopenharmony_ci
43162306a36Sopenharmony_ci        digest_cnt = struct.unpack('>I', rsp[:4])[0]
43262306a36Sopenharmony_ci        if digest_cnt == 0:
43362306a36Sopenharmony_ci            return None
43462306a36Sopenharmony_ci        rsp = rsp[6:]
43562306a36Sopenharmony_ci
43662306a36Sopenharmony_ci        return rsp
43762306a36Sopenharmony_ci
43862306a36Sopenharmony_ci    def extend_pcr(self, i, dig, bank_alg = TPM2_ALG_SHA1):
43962306a36Sopenharmony_ci        ds = get_digest_size(bank_alg)
44062306a36Sopenharmony_ci        assert(ds == len(dig))
44162306a36Sopenharmony_ci
44262306a36Sopenharmony_ci        auth_cmd = AuthCommand()
44362306a36Sopenharmony_ci
44462306a36Sopenharmony_ci        fmt = '>HII I I%us IH%us' % (len(auth_cmd), ds)
44562306a36Sopenharmony_ci        cmd = struct.pack(
44662306a36Sopenharmony_ci            fmt,
44762306a36Sopenharmony_ci            TPM2_ST_SESSIONS,
44862306a36Sopenharmony_ci            struct.calcsize(fmt),
44962306a36Sopenharmony_ci            TPM2_CC_PCR_EXTEND,
45062306a36Sopenharmony_ci            i,
45162306a36Sopenharmony_ci            len(auth_cmd),
45262306a36Sopenharmony_ci            bytes(auth_cmd),
45362306a36Sopenharmony_ci            1, bank_alg, dig)
45462306a36Sopenharmony_ci
45562306a36Sopenharmony_ci        self.send_cmd(cmd)
45662306a36Sopenharmony_ci
45762306a36Sopenharmony_ci    def start_auth_session(self, session_type, name_alg = TPM2_ALG_SHA1):
45862306a36Sopenharmony_ci        fmt = '>HII IIH16sHBHH'
45962306a36Sopenharmony_ci        cmd = struct.pack(fmt,
46062306a36Sopenharmony_ci                          TPM2_ST_NO_SESSIONS,
46162306a36Sopenharmony_ci                          struct.calcsize(fmt),
46262306a36Sopenharmony_ci                          TPM2_CC_START_AUTH_SESSION,
46362306a36Sopenharmony_ci                          TPM2_RH_NULL,
46462306a36Sopenharmony_ci                          TPM2_RH_NULL,
46562306a36Sopenharmony_ci                          16,
46662306a36Sopenharmony_ci                          ('\0' * 16).encode(),
46762306a36Sopenharmony_ci                          0,
46862306a36Sopenharmony_ci                          session_type,
46962306a36Sopenharmony_ci                          TPM2_ALG_NULL,
47062306a36Sopenharmony_ci                          name_alg)
47162306a36Sopenharmony_ci
47262306a36Sopenharmony_ci        return struct.unpack('>I', self.send_cmd(cmd)[10:14])[0]
47362306a36Sopenharmony_ci
47462306a36Sopenharmony_ci    def __calc_pcr_digest(self, pcrs, bank_alg = TPM2_ALG_SHA1,
47562306a36Sopenharmony_ci                          digest_alg = TPM2_ALG_SHA1):
47662306a36Sopenharmony_ci        x = []
47762306a36Sopenharmony_ci        f = get_hash_function(digest_alg)
47862306a36Sopenharmony_ci
47962306a36Sopenharmony_ci        for i in pcrs:
48062306a36Sopenharmony_ci            pcr = self.read_pcr(i, bank_alg)
48162306a36Sopenharmony_ci            if pcr is None:
48262306a36Sopenharmony_ci                return None
48362306a36Sopenharmony_ci            x += pcr
48462306a36Sopenharmony_ci
48562306a36Sopenharmony_ci        return f(bytearray(x)).digest()
48662306a36Sopenharmony_ci
48762306a36Sopenharmony_ci    def policy_pcr(self, handle, pcrs, bank_alg = TPM2_ALG_SHA1,
48862306a36Sopenharmony_ci                   name_alg = TPM2_ALG_SHA1):
48962306a36Sopenharmony_ci        ds = get_digest_size(name_alg)
49062306a36Sopenharmony_ci        dig = self.__calc_pcr_digest(pcrs, bank_alg, name_alg)
49162306a36Sopenharmony_ci        if not dig:
49262306a36Sopenharmony_ci            raise UnknownPCRBankError(bank_alg)
49362306a36Sopenharmony_ci
49462306a36Sopenharmony_ci        pcrsel_len = max((max(pcrs) >> 3) + 1, 3)
49562306a36Sopenharmony_ci        pcrsel = [0] * pcrsel_len
49662306a36Sopenharmony_ci        for i in pcrs:
49762306a36Sopenharmony_ci            pcrsel[i >> 3] |= 1 << (i & 7)
49862306a36Sopenharmony_ci        pcrsel = ''.join(map(chr, pcrsel)).encode()
49962306a36Sopenharmony_ci
50062306a36Sopenharmony_ci        fmt = '>HII IH%usIHB3s' % ds
50162306a36Sopenharmony_ci        cmd = struct.pack(fmt,
50262306a36Sopenharmony_ci                          TPM2_ST_NO_SESSIONS,
50362306a36Sopenharmony_ci                          struct.calcsize(fmt),
50462306a36Sopenharmony_ci                          TPM2_CC_POLICY_PCR,
50562306a36Sopenharmony_ci                          handle,
50662306a36Sopenharmony_ci                          len(dig),
50762306a36Sopenharmony_ci                          bytes(dig),
50862306a36Sopenharmony_ci                          1,
50962306a36Sopenharmony_ci                          bank_alg,
51062306a36Sopenharmony_ci                          pcrsel_len, pcrsel)
51162306a36Sopenharmony_ci
51262306a36Sopenharmony_ci        self.send_cmd(cmd)
51362306a36Sopenharmony_ci
51462306a36Sopenharmony_ci    def policy_password(self, handle):
51562306a36Sopenharmony_ci        fmt = '>HII I'
51662306a36Sopenharmony_ci        cmd = struct.pack(fmt,
51762306a36Sopenharmony_ci                          TPM2_ST_NO_SESSIONS,
51862306a36Sopenharmony_ci                          struct.calcsize(fmt),
51962306a36Sopenharmony_ci                          TPM2_CC_POLICY_PASSWORD,
52062306a36Sopenharmony_ci                          handle)
52162306a36Sopenharmony_ci
52262306a36Sopenharmony_ci        self.send_cmd(cmd)
52362306a36Sopenharmony_ci
52462306a36Sopenharmony_ci    def get_policy_digest(self, handle):
52562306a36Sopenharmony_ci        fmt = '>HII I'
52662306a36Sopenharmony_ci        cmd = struct.pack(fmt,
52762306a36Sopenharmony_ci                          TPM2_ST_NO_SESSIONS,
52862306a36Sopenharmony_ci                          struct.calcsize(fmt),
52962306a36Sopenharmony_ci                          TPM2_CC_POLICY_GET_DIGEST,
53062306a36Sopenharmony_ci                          handle)
53162306a36Sopenharmony_ci
53262306a36Sopenharmony_ci        return self.send_cmd(cmd)[12:]
53362306a36Sopenharmony_ci
53462306a36Sopenharmony_ci    def flush_context(self, handle):
53562306a36Sopenharmony_ci        fmt = '>HIII'
53662306a36Sopenharmony_ci        cmd = struct.pack(fmt,
53762306a36Sopenharmony_ci                          TPM2_ST_NO_SESSIONS,
53862306a36Sopenharmony_ci                          struct.calcsize(fmt),
53962306a36Sopenharmony_ci                          TPM2_CC_FLUSH_CONTEXT,
54062306a36Sopenharmony_ci                          handle)
54162306a36Sopenharmony_ci
54262306a36Sopenharmony_ci        self.send_cmd(cmd)
54362306a36Sopenharmony_ci
54462306a36Sopenharmony_ci    def create_root_key(self, auth_value = bytes()):
54562306a36Sopenharmony_ci        attributes = \
54662306a36Sopenharmony_ci            Public.FIXED_TPM | \
54762306a36Sopenharmony_ci            Public.FIXED_PARENT | \
54862306a36Sopenharmony_ci            Public.SENSITIVE_DATA_ORIGIN | \
54962306a36Sopenharmony_ci            Public.USER_WITH_AUTH | \
55062306a36Sopenharmony_ci            Public.RESTRICTED | \
55162306a36Sopenharmony_ci            Public.DECRYPT
55262306a36Sopenharmony_ci
55362306a36Sopenharmony_ci        auth_cmd = AuthCommand()
55462306a36Sopenharmony_ci        sensitive = SensitiveCreate(user_auth=auth_value)
55562306a36Sopenharmony_ci
55662306a36Sopenharmony_ci        public_parms = struct.pack(
55762306a36Sopenharmony_ci            '>HHHHHI',
55862306a36Sopenharmony_ci            TPM2_ALG_AES,
55962306a36Sopenharmony_ci            128,
56062306a36Sopenharmony_ci            TPM2_ALG_CFB,
56162306a36Sopenharmony_ci            TPM2_ALG_NULL,
56262306a36Sopenharmony_ci            2048,
56362306a36Sopenharmony_ci            0)
56462306a36Sopenharmony_ci
56562306a36Sopenharmony_ci        public = Public(
56662306a36Sopenharmony_ci            object_type=TPM2_ALG_RSA,
56762306a36Sopenharmony_ci            name_alg=TPM2_ALG_SHA1,
56862306a36Sopenharmony_ci            object_attributes=attributes,
56962306a36Sopenharmony_ci            parameters=public_parms)
57062306a36Sopenharmony_ci
57162306a36Sopenharmony_ci        fmt = '>HIII I%us H%us H%us HI' % \
57262306a36Sopenharmony_ci            (len(auth_cmd), len(sensitive), len(public))
57362306a36Sopenharmony_ci        cmd = struct.pack(
57462306a36Sopenharmony_ci            fmt,
57562306a36Sopenharmony_ci            TPM2_ST_SESSIONS,
57662306a36Sopenharmony_ci            struct.calcsize(fmt),
57762306a36Sopenharmony_ci            TPM2_CC_CREATE_PRIMARY,
57862306a36Sopenharmony_ci            TPM2_RH_OWNER,
57962306a36Sopenharmony_ci            len(auth_cmd),
58062306a36Sopenharmony_ci            bytes(auth_cmd),
58162306a36Sopenharmony_ci            len(sensitive),
58262306a36Sopenharmony_ci            bytes(sensitive),
58362306a36Sopenharmony_ci            len(public),
58462306a36Sopenharmony_ci            bytes(public),
58562306a36Sopenharmony_ci            0, 0)
58662306a36Sopenharmony_ci
58762306a36Sopenharmony_ci        return struct.unpack('>I', self.send_cmd(cmd)[10:14])[0]
58862306a36Sopenharmony_ci
58962306a36Sopenharmony_ci    def seal(self, parent_key, data, auth_value, policy_dig,
59062306a36Sopenharmony_ci             name_alg = TPM2_ALG_SHA1):
59162306a36Sopenharmony_ci        ds = get_digest_size(name_alg)
59262306a36Sopenharmony_ci        assert(not policy_dig or ds == len(policy_dig))
59362306a36Sopenharmony_ci
59462306a36Sopenharmony_ci        attributes = 0
59562306a36Sopenharmony_ci        if not policy_dig:
59662306a36Sopenharmony_ci            attributes |= Public.USER_WITH_AUTH
59762306a36Sopenharmony_ci            policy_dig = bytes()
59862306a36Sopenharmony_ci
59962306a36Sopenharmony_ci        auth_cmd =  AuthCommand()
60062306a36Sopenharmony_ci        sensitive = SensitiveCreate(user_auth=auth_value, data=data)
60162306a36Sopenharmony_ci
60262306a36Sopenharmony_ci        public = Public(
60362306a36Sopenharmony_ci            object_type=TPM2_ALG_KEYEDHASH,
60462306a36Sopenharmony_ci            name_alg=name_alg,
60562306a36Sopenharmony_ci            object_attributes=attributes,
60662306a36Sopenharmony_ci            auth_policy=policy_dig,
60762306a36Sopenharmony_ci            parameters=struct.pack('>H', TPM2_ALG_NULL))
60862306a36Sopenharmony_ci
60962306a36Sopenharmony_ci        fmt = '>HIII I%us H%us H%us HI' % \
61062306a36Sopenharmony_ci            (len(auth_cmd), len(sensitive), len(public))
61162306a36Sopenharmony_ci        cmd = struct.pack(
61262306a36Sopenharmony_ci            fmt,
61362306a36Sopenharmony_ci            TPM2_ST_SESSIONS,
61462306a36Sopenharmony_ci            struct.calcsize(fmt),
61562306a36Sopenharmony_ci            TPM2_CC_CREATE,
61662306a36Sopenharmony_ci            parent_key,
61762306a36Sopenharmony_ci            len(auth_cmd),
61862306a36Sopenharmony_ci            bytes(auth_cmd),
61962306a36Sopenharmony_ci            len(sensitive),
62062306a36Sopenharmony_ci            bytes(sensitive),
62162306a36Sopenharmony_ci            len(public),
62262306a36Sopenharmony_ci            bytes(public),
62362306a36Sopenharmony_ci            0, 0)
62462306a36Sopenharmony_ci
62562306a36Sopenharmony_ci        rsp = self.send_cmd(cmd)
62662306a36Sopenharmony_ci
62762306a36Sopenharmony_ci        return rsp[14:]
62862306a36Sopenharmony_ci
62962306a36Sopenharmony_ci    def unseal(self, parent_key, blob, auth_value, policy_handle):
63062306a36Sopenharmony_ci        private_len = struct.unpack('>H', blob[0:2])[0]
63162306a36Sopenharmony_ci        public_start = private_len + 2
63262306a36Sopenharmony_ci        public_len = struct.unpack('>H', blob[public_start:public_start + 2])[0]
63362306a36Sopenharmony_ci        blob = blob[:private_len + public_len + 4]
63462306a36Sopenharmony_ci
63562306a36Sopenharmony_ci        auth_cmd = AuthCommand()
63662306a36Sopenharmony_ci
63762306a36Sopenharmony_ci        fmt = '>HII I I%us %us' % (len(auth_cmd), len(blob))
63862306a36Sopenharmony_ci        cmd = struct.pack(
63962306a36Sopenharmony_ci            fmt,
64062306a36Sopenharmony_ci            TPM2_ST_SESSIONS,
64162306a36Sopenharmony_ci            struct.calcsize(fmt),
64262306a36Sopenharmony_ci            TPM2_CC_LOAD,
64362306a36Sopenharmony_ci            parent_key,
64462306a36Sopenharmony_ci            len(auth_cmd),
64562306a36Sopenharmony_ci            bytes(auth_cmd),
64662306a36Sopenharmony_ci            blob)
64762306a36Sopenharmony_ci
64862306a36Sopenharmony_ci        data_handle = struct.unpack('>I', self.send_cmd(cmd)[10:14])[0]
64962306a36Sopenharmony_ci
65062306a36Sopenharmony_ci        if policy_handle:
65162306a36Sopenharmony_ci            auth_cmd = AuthCommand(session_handle=policy_handle, hmac=auth_value)
65262306a36Sopenharmony_ci        else:
65362306a36Sopenharmony_ci            auth_cmd = AuthCommand(hmac=auth_value)
65462306a36Sopenharmony_ci
65562306a36Sopenharmony_ci        fmt = '>HII I I%us' % (len(auth_cmd))
65662306a36Sopenharmony_ci        cmd = struct.pack(
65762306a36Sopenharmony_ci            fmt,
65862306a36Sopenharmony_ci            TPM2_ST_SESSIONS,
65962306a36Sopenharmony_ci            struct.calcsize(fmt),
66062306a36Sopenharmony_ci            TPM2_CC_UNSEAL,
66162306a36Sopenharmony_ci            data_handle,
66262306a36Sopenharmony_ci            len(auth_cmd),
66362306a36Sopenharmony_ci            bytes(auth_cmd))
66462306a36Sopenharmony_ci
66562306a36Sopenharmony_ci        try:
66662306a36Sopenharmony_ci            rsp = self.send_cmd(cmd)
66762306a36Sopenharmony_ci        finally:
66862306a36Sopenharmony_ci            self.flush_context(data_handle)
66962306a36Sopenharmony_ci
67062306a36Sopenharmony_ci        data_len = struct.unpack('>I', rsp[10:14])[0] - 2
67162306a36Sopenharmony_ci
67262306a36Sopenharmony_ci        return rsp[16:16 + data_len]
67362306a36Sopenharmony_ci
67462306a36Sopenharmony_ci    def reset_da_lock(self):
67562306a36Sopenharmony_ci        auth_cmd = AuthCommand()
67662306a36Sopenharmony_ci
67762306a36Sopenharmony_ci        fmt = '>HII I I%us' % (len(auth_cmd))
67862306a36Sopenharmony_ci        cmd = struct.pack(
67962306a36Sopenharmony_ci            fmt,
68062306a36Sopenharmony_ci            TPM2_ST_SESSIONS,
68162306a36Sopenharmony_ci            struct.calcsize(fmt),
68262306a36Sopenharmony_ci            TPM2_CC_DICTIONARY_ATTACK_LOCK_RESET,
68362306a36Sopenharmony_ci            TPM2_RH_LOCKOUT,
68462306a36Sopenharmony_ci            len(auth_cmd),
68562306a36Sopenharmony_ci            bytes(auth_cmd))
68662306a36Sopenharmony_ci
68762306a36Sopenharmony_ci        self.send_cmd(cmd)
68862306a36Sopenharmony_ci
68962306a36Sopenharmony_ci    def __get_cap_cnt(self, cap, pt, cnt):
69062306a36Sopenharmony_ci        handles = []
69162306a36Sopenharmony_ci        fmt = '>HII III'
69262306a36Sopenharmony_ci
69362306a36Sopenharmony_ci        cmd = struct.pack(fmt,
69462306a36Sopenharmony_ci                          TPM2_ST_NO_SESSIONS,
69562306a36Sopenharmony_ci                          struct.calcsize(fmt),
69662306a36Sopenharmony_ci                          TPM2_CC_GET_CAPABILITY,
69762306a36Sopenharmony_ci                          cap, pt, cnt)
69862306a36Sopenharmony_ci
69962306a36Sopenharmony_ci        rsp = self.send_cmd(cmd)[10:]
70062306a36Sopenharmony_ci        more_data, cap, cnt = struct.unpack('>BII', rsp[:9])
70162306a36Sopenharmony_ci        rsp = rsp[9:]
70262306a36Sopenharmony_ci
70362306a36Sopenharmony_ci        for i in range(0, cnt):
70462306a36Sopenharmony_ci            handle = struct.unpack('>I', rsp[:4])[0]
70562306a36Sopenharmony_ci            handles.append(handle)
70662306a36Sopenharmony_ci            rsp = rsp[4:]
70762306a36Sopenharmony_ci
70862306a36Sopenharmony_ci        return handles, more_data
70962306a36Sopenharmony_ci
71062306a36Sopenharmony_ci    def get_cap(self, cap, pt):
71162306a36Sopenharmony_ci        handles = []
71262306a36Sopenharmony_ci
71362306a36Sopenharmony_ci        more_data = True
71462306a36Sopenharmony_ci        while more_data:
71562306a36Sopenharmony_ci            next_handles, more_data = self.__get_cap_cnt(cap, pt, 1)
71662306a36Sopenharmony_ci            handles += next_handles
71762306a36Sopenharmony_ci            pt += 1
71862306a36Sopenharmony_ci
71962306a36Sopenharmony_ci        return handles
72062306a36Sopenharmony_ci
72162306a36Sopenharmony_ci    def get_cap_pcrs(self):
72262306a36Sopenharmony_ci        pcr_banks = {}
72362306a36Sopenharmony_ci
72462306a36Sopenharmony_ci        fmt = '>HII III'
72562306a36Sopenharmony_ci
72662306a36Sopenharmony_ci        cmd = struct.pack(fmt,
72762306a36Sopenharmony_ci                          TPM2_ST_NO_SESSIONS,
72862306a36Sopenharmony_ci                          struct.calcsize(fmt),
72962306a36Sopenharmony_ci                          TPM2_CC_GET_CAPABILITY,
73062306a36Sopenharmony_ci                          TPM2_CAP_PCRS, 0, 1)
73162306a36Sopenharmony_ci        rsp = self.send_cmd(cmd)[10:]
73262306a36Sopenharmony_ci        _, _, cnt = struct.unpack('>BII', rsp[:9])
73362306a36Sopenharmony_ci        rsp = rsp[9:]
73462306a36Sopenharmony_ci
73562306a36Sopenharmony_ci        # items are TPMS_PCR_SELECTION's
73662306a36Sopenharmony_ci        for i in range(0, cnt):
73762306a36Sopenharmony_ci              hash, sizeOfSelect = struct.unpack('>HB', rsp[:3])
73862306a36Sopenharmony_ci              rsp = rsp[3:]
73962306a36Sopenharmony_ci
74062306a36Sopenharmony_ci              pcrSelect = 0
74162306a36Sopenharmony_ci              if sizeOfSelect > 0:
74262306a36Sopenharmony_ci                  pcrSelect, = struct.unpack('%ds' % sizeOfSelect,
74362306a36Sopenharmony_ci                                             rsp[:sizeOfSelect])
74462306a36Sopenharmony_ci                  rsp = rsp[sizeOfSelect:]
74562306a36Sopenharmony_ci                  pcrSelect = int.from_bytes(pcrSelect, byteorder='big')
74662306a36Sopenharmony_ci
74762306a36Sopenharmony_ci              pcr_banks[hash] = pcrSelect
74862306a36Sopenharmony_ci
74962306a36Sopenharmony_ci        return pcr_banks
750