xref: /aosp_15_r20/external/pigweed/pw_software_update/py/pw_software_update/keys.py (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1# Copyright 2021 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Facilities for keys generation, importing, signing and verification.
15
16IMPORTANT: THESE FACILITIES ARE FOR LOCAL NON-PRODUCTION USE ONLY!!
17
18These are not suited for production use because:
19
201. The private keys are not generated without ANY supervision or authorization.
212. The private keys are not stored securely.
223. The underlying crypto library is not audited.
23"""
24
25import argparse
26import hashlib
27from pathlib import Path
28
29from cryptography.hazmat.primitives import hashes
30from cryptography.hazmat.primitives.asymmetric import ec
31from cryptography.hazmat.primitives.asymmetric.utils import (
32    decode_dss_signature,
33    encode_dss_signature,
34)
35from cryptography.hazmat.primitives.serialization import (
36    Encoding,
37    NoEncryption,
38    PrivateFormat,
39    PublicFormat,
40    load_pem_private_key,
41    load_pem_public_key,
42)
43
44from pw_software_update.tuf_pb2 import (
45    Key,
46    KeyMapping,
47    KeyScheme,
48    KeyType,
49    Signature,
50)
51
52
53def parse_args():
54    """Parse CLI arguments."""
55    parser = argparse.ArgumentParser(description=__doc__)
56    parser.add_argument(
57        '-o',
58        '--out',
59        type=Path,
60        required=True,
61        help='Output path for the generated key',
62    )
63    return parser.parse_args()
64
65
66def gen_ecdsa_keypair(out: Path) -> None:
67    """Generates and writes to disk a NIST-P256 EC key pair.
68
69    Args:
70      out: The path to write the private key to. The public key is written
71        to the same path as the private key using the suffix '.pub'.
72    """
73    private_key = ec.generate_private_key(ec.SECP256R1())
74    public_key = private_key.public_key()
75    private_pem = private_key.private_bytes(
76        encoding=Encoding.PEM,
77        format=PrivateFormat.PKCS8,
78        encryption_algorithm=NoEncryption(),
79    )
80    public_pem = public_key.public_bytes(
81        encoding=Encoding.PEM, format=PublicFormat.SubjectPublicKeyInfo
82    )
83
84    out.write_bytes(private_pem)
85    public_out = out.parent / f'{out.name}.pub'
86    public_out.write_bytes(public_pem)
87
88
89def gen_key_id(key: Key) -> bytes:
90    """Computes the key ID of a Key object."""
91    sha = hashlib.sha256()
92    sha.update(key.key_type.to_bytes(1, 'big'))
93    sha.update(key.scheme.to_bytes(1, 'big'))
94    sha.update(key.keyval)
95    return sha.digest()
96
97
98def import_ecdsa_public_key(pem: bytes) -> KeyMapping:
99    """Imports an EC NIST-P256 public key in pem format."""
100    ec_key = load_pem_public_key(pem)
101
102    if not isinstance(ec_key, ec.EllipticCurvePublicKey):
103        raise TypeError(
104            f'Not an elliptic curve public key type: {type(ec_key)}.'
105            'Try generate a key with gen_ecdsa_keypair()?'
106        )
107
108    # pylint: disable=no-member
109    if not (ec_key.curve.name == 'secp256r1' and ec_key.key_size == 256):
110        raise TypeError(
111            f'Unsupported curve: {ec_key.curve.name}.'
112            'Try generate a key with gen_ecdsa_keypair()?'
113        )
114    # pylint: enable=no-member
115
116    tuf_key = Key(
117        key_type=KeyType.ECDSA_SHA2_NISTP256,
118        scheme=KeyScheme.ECDSA_SHA2_NISTP256_SCHEME,
119        keyval=ec_key.public_bytes(
120            Encoding.X962, PublicFormat.UncompressedPoint
121        ),
122    )
123    return KeyMapping(key_id=gen_key_id(tuf_key), key=tuf_key)
124
125
126def create_ecdsa_signature(data: bytes, key: bytes) -> Signature:
127    """Creates an ECDSA-SHA2-NISTP256 signature."""
128    ec_key = load_pem_private_key(key, password=None)
129    if not isinstance(ec_key, ec.EllipticCurvePrivateKey):
130        raise TypeError(
131            f'Not an elliptic curve private key: {type(ec_key)}.'
132            'Try generate a key with gen_ecdsa_keypair()?'
133        )
134
135    tuf_key = Key(
136        key_type=KeyType.ECDSA_SHA2_NISTP256,
137        scheme=KeyScheme.ECDSA_SHA2_NISTP256_SCHEME,
138        keyval=ec_key.public_key().public_bytes(
139            Encoding.X962, PublicFormat.UncompressedPoint
140        ),
141    )
142
143    der_signature = ec_key.sign(
144        data, ec.ECDSA(hashes.SHA256())
145    )  # pylint: disable=no-value-for-parameter
146    int_r, int_s = decode_dss_signature(der_signature)
147    sig_bytes = int_r.to_bytes(32, 'big') + int_s.to_bytes(32, 'big')
148
149    return Signature(key_id=gen_key_id(tuf_key), sig=sig_bytes)
150
151
152def verify_ecdsa_signature(sig: bytes, data: bytes, key: Key) -> bool:
153    """Verifies an ECDSA-SHA2-NISTP256 signature with a given public key.
154
155    Args:
156      sig: the ECDSA signature as raw bytes (r||s).
157      data: the message as plain text.
158      key: the ECDSA-NISTP256 public key.
159
160    Returns:
161      True if the signature is verified. False otherwise.
162    """
163    ec_key = ec.EllipticCurvePublicKey.from_encoded_point(
164        ec.SECP256R1(), key.keyval
165    )
166    try:
167        dss_sig = encode_dss_signature(
168            int.from_bytes(sig[:32], 'big'), int.from_bytes(sig[-32:], 'big')
169        )
170        ec_key.verify(dss_sig, data, ec.ECDSA(hashes.SHA256()))
171    except:  # pylint: disable=bare-except
172        return False
173
174    return True
175
176
177def main(out: Path) -> None:
178    """Generates and writes to disk key pairs for development use."""
179
180    # Currently only supports the "ecdsa-sha2-nistp256" key scheme.
181    #
182    # TODO(alizhang): Add support for "rsassa-pss-sha256" and "ed25519" key
183    # schemes.
184    gen_ecdsa_keypair(out)
185
186
187if __name__ == '__main__':
188    main(**vars(parse_args()))
189