1# Copyright 2018 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import json
6import logging
7
8from subprocess import CalledProcessError
9
10
11class PinWeaverNotAvailableError(CalledProcessError):
12    """This exception is thrown when pinweaver_client reports that the PinWeaver
13    feature is not available.
14    """
15
16    def __init__(self, *args, **kwargs):
17        super(PinWeaverNotAvailableError, self).__init__(*args, **kwargs)
18
19
20def __check_pinweaver_client_present(client, message):
21    cmd = 'which pinweaver_client'
22    run = client.run('which pinweaver_client', ignore_status=True)
23    if run.exit_status != 0:  # pinweaver_client isn't present.
24        raise PinWeaverNotAvailableError(run.exit_status, cmd, message);
25
26def __execute_for_dict(client, *args, **kwargs):
27    """Executes a command with the specified args and parses stdout as JSON
28    based on the expected output of pinweaver_client.
29    """
30    __check_pinweaver_client_present(client, args[0])
31
32    result = {}
33    stack = [result]
34    if 'ignore_status' not in kwargs:
35        kwargs['ignore_status'] = True
36    run = client.run(*args, **kwargs)
37    if run.exit_status == 2:  # EXIT_PINWEAVER_NOT_SUPPORTED
38        raise PinWeaverNotAvailableError(run.exit_status, args[0]);
39    logging.debug(args)
40    logging.info(run.stderr)
41    logging.debug(run.stdout)
42    return json.loads(run.stdout)
43
44
45def ResetTree(client, bits_per_level, height):
46    """Returns a dictionary with keys result_code and root_hash.
47
48    @param client: client object to run commands on.
49    """
50    return __execute_for_dict(client, 'pinweaver_client resettree %d %d' %
51                              (bits_per_level, height))
52
53
54def InsertLeaf(client, label, auxilary_hashes, low_entropy_secret,
55               high_entropy_secret, reset_secret, delay_schedule):
56    """Returns a dictionary with keys result_code, root_hash, cred_metadata,
57    and mac.
58
59    @param client: client object to run commands on.
60    """
61    return __execute_for_dict(
62        client, 'pinweaver_client insert %d %s %s %s %s %s' %
63                (label, auxilary_hashes, low_entropy_secret,
64                 high_entropy_secret, reset_secret, delay_schedule))
65
66
67def RemoveLeaf(client, label, auxilary_hashes, mac):
68    """Returns a dictionary with keys result_code and root_hash.
69
70    @param client: client object to run commands on.
71    """
72    return __execute_for_dict(
73        client, 'pinweaver_client remove %d %s %s' %
74                (label, auxilary_hashes, mac))
75
76
77def TryAuth(client, auxilary_hashes, low_entropy_secret, cred_metadata):
78    """Returns a dictionary with keys result_code, root_hash, cred_metadata,
79    mac, and he_secret.
80
81    @param client: client object to run commands on.
82    """
83    return __execute_for_dict(
84        client, 'pinweaver_client auth %s %s %s' %
85                (auxilary_hashes, low_entropy_secret, cred_metadata))
86
87
88def ResetAuth(client, auxilary_hashes, reset_secret, cred_metadata):
89    """Returns a dictionary with keys result_code, root_hash, cred_metadata,
90    mac, and he_secret.
91
92    @param client: client object to run commands on.
93    """
94    return __execute_for_dict(
95        client, 'pinweaver_client resetleaf %s %s %s' %
96                (auxilary_hashes, reset_secret, cred_metadata))
97
98
99def GetLog(client, root=None):
100    """Returns a dictionary with keys result_code, root_hash, and a list of
101    entry[#] sub dictionaries for each log entry.
102
103    @param client: client object to run commands on.
104    @param root: root hash of the log entry to search for.
105    """
106    if root is None:
107        root = ('0' * 64)
108
109    return __execute_for_dict(client, 'pinweaver_client getlog %s' % (root))
110
111
112def LogReplay(client, auxilary_hashes, log_root, cred_metadata):
113    """Returns a dictionary with keys result_code, root_hash, cred_metadata,
114    and mac.
115
116    @param client: client object to run commands on.
117    """
118    return __execute_for_dict(
119        client, 'pinweaver_client replay %d %s %s %s' %
120                (auxilary_hashes, log_root, cred_metadata))
121
122
123def SelfTest(client):
124    """Returns True if the test succeeded.
125
126    @param client: client object to run commands on.
127    """
128    cmd = 'pinweaver_client selftest'
129    __check_pinweaver_client_present(client, cmd)
130
131    run = client.run(cmd)
132    if run.exit_status == -2:
133        raise PinWeaverNotAvailableError(run.exit_status, cmd);
134    output = run.stdout
135    return "Success!" in output
136