|
14 | 14 | # See the License for the specific language governing permissions and
|
15 | 15 | # limitations under the License.
|
16 | 16 |
|
17 |
| -import click |
18 |
| -import getpass |
19 |
| -from imgtool import keys |
20 |
| -from imgtool import image |
21 |
| -from imgtool.version import decode_version |
22 |
| - |
23 |
| - |
24 |
| -def gen_rsa2048(keyfile, passwd): |
25 |
| - keys.RSA2048.generate().export_private(path=keyfile, passwd=passwd) |
26 |
| - |
27 |
| - |
28 |
| -def gen_ecdsa_p256(keyfile, passwd): |
29 |
| - keys.ECDSA256P1.generate().export_private(keyfile, passwd=passwd) |
30 |
| - |
31 |
| - |
32 |
| -def gen_ecdsa_p224(keyfile, passwd): |
33 |
| - print("TODO: p-224 not yet implemented") |
34 |
| - |
35 |
| - |
36 |
| -valid_langs = ['c', 'rust'] |
37 |
| -keygens = { |
38 |
| - 'rsa-2048': gen_rsa2048, |
39 |
| - 'ecdsa-p256': gen_ecdsa_p256, |
40 |
| - 'ecdsa-p224': gen_ecdsa_p224, |
41 |
| -} |
42 |
| - |
43 |
| - |
44 |
| -def load_key(keyfile): |
45 |
| - # TODO: better handling of invalid pass-phrase |
46 |
| - key = keys.load(keyfile) |
47 |
| - if key is not None: |
48 |
| - return key |
49 |
| - passwd = getpass.getpass("Enter key passphrase: ").encode('utf-8') |
50 |
| - return keys.load(keyfile, passwd) |
51 |
| - |
52 |
| - |
53 |
| -def get_password(): |
54 |
| - while True: |
55 |
| - passwd = getpass.getpass("Enter key passphrase: ") |
56 |
| - passwd2 = getpass.getpass("Reenter passphrase: ") |
57 |
| - if passwd == passwd2: |
58 |
| - break |
59 |
| - print("Passwords do not match, try again") |
60 |
| - |
61 |
| - # Password must be bytes, always use UTF-8 for consistent |
62 |
| - # encoding. |
63 |
| - return passwd.encode('utf-8') |
64 |
| - |
65 |
| - |
66 |
| -@click.option('-p', '--password', is_flag=True, |
67 |
| - help='Prompt for password to protect key') |
68 |
| -@click.option('-t', '--type', metavar='type', required=True, |
69 |
| - type=click.Choice(keygens.keys())) |
70 |
| -@click.option('-k', '--key', metavar='filename', required=True) |
71 |
| -@click.command(help='Generate pub/private keypair') |
72 |
| -def keygen(type, key, password): |
73 |
| - password = get_password() if password else None |
74 |
| - keygens[type](key, password) |
75 |
| - |
76 |
| - |
77 |
| -@click.option('-l', '--lang', metavar='lang', default=valid_langs[0], |
78 |
| - type=click.Choice(valid_langs)) |
79 |
| -@click.option('-k', '--key', metavar='filename', required=True) |
80 |
| -@click.command(help='Get public key from keypair') |
81 |
| -def getpub(key, lang): |
82 |
| - key = load_key(key) |
83 |
| - if key is None: |
84 |
| - print("Invalid passphrase") |
85 |
| - elif lang == 'c': |
86 |
| - key.emit_c() |
87 |
| - elif lang == 'rust': |
88 |
| - key.emit_rust() |
89 |
| - else: |
90 |
| - raise ValueError("BUG: should never get here!") |
91 |
| - |
92 |
| - |
93 |
| -def validate_version(ctx, param, value): |
94 |
| - try: |
95 |
| - decode_version(value) |
96 |
| - return value |
97 |
| - except ValueError as e: |
98 |
| - raise click.BadParameter("{}".format(e)) |
99 |
| - |
100 |
| - |
101 |
| -def validate_header_size(ctx, param, value): |
102 |
| - min_hdr_size = image.IMAGE_HEADER_SIZE |
103 |
| - if value < min_hdr_size: |
104 |
| - raise click.BadParameter( |
105 |
| - "Minimum value for -H/--header-size is {}".format(min_hdr_size)) |
106 |
| - return value |
107 |
| - |
108 |
| - |
109 |
| -class BasedIntParamType(click.ParamType): |
110 |
| - name = 'integer' |
111 |
| - |
112 |
| - def convert(self, value, param, ctx): |
113 |
| - try: |
114 |
| - if value[:2].lower() == '0x': |
115 |
| - return int(value[2:], 16) |
116 |
| - elif value[:1] == '0': |
117 |
| - return int(value, 8) |
118 |
| - return int(value, 10) |
119 |
| - except ValueError: |
120 |
| - self.fail('%s is not a valid integer' % value, param, ctx) |
121 |
| - |
122 |
| - |
123 |
| -@click.argument('outfile') |
124 |
| -@click.argument('infile') |
125 |
| -@click.option('-E', '--encrypt', metavar='filename', |
126 |
| - help='Encrypt image using the provided public key') |
127 |
| -@click.option('-e', '--endian', type=click.Choice(['little', 'big']), |
128 |
| - default='little', help="Select little or big endian") |
129 |
| -@click.option('--overwrite-only', default=False, is_flag=True, |
130 |
| - help='Use overwrite-only instead of swap upgrades') |
131 |
| -@click.option('-M', '--max-sectors', type=int, |
132 |
| - help='When padding allow for this amount of sectors (defaults to 128)') |
133 |
| -@click.option('--pad', default=False, is_flag=True, |
134 |
| - help='Pad image to --slot-size bytes, adding trailer magic') |
135 |
| -@click.option('-S', '--slot-size', type=BasedIntParamType(), required=True, |
136 |
| - help='Size of the slot where the image will be written') |
137 |
| -@click.option('--pad-header', default=False, is_flag=True, |
138 |
| - help='Add --header-size zeroed bytes at the beginning of the image') |
139 |
| -@click.option('-H', '--header-size', callback=validate_header_size, |
140 |
| - type=BasedIntParamType(), required=True) |
141 |
| -@click.option('-v', '--version', callback=validate_version, required=True) |
142 |
| -@click.option('--align', type=click.Choice(['1', '2', '4', '8']), |
143 |
| - required=True) |
144 |
| -@click.option('-k', '--key', metavar='filename') |
145 |
| -@click.command(help='Create a signed or unsigned image') |
146 |
| -def sign(key, align, version, header_size, pad_header, slot_size, pad, |
147 |
| - max_sectors, overwrite_only, endian, encrypt, infile, outfile): |
148 |
| - img = image.Image.load(infile, version=decode_version(version), |
149 |
| - header_size=header_size, pad_header=pad_header, |
150 |
| - pad=pad, align=int(align), slot_size=slot_size, |
151 |
| - max_sectors=max_sectors, |
152 |
| - overwrite_only=overwrite_only, |
153 |
| - endian=endian) |
154 |
| - key = load_key(key) if key else None |
155 |
| - enckey = load_key(encrypt) if encrypt else None |
156 |
| - if enckey: |
157 |
| - if not isinstance(enckey, (keys.RSA2048, keys.RSA2048Public)): |
158 |
| - raise Exception("Encryption only available with RSA") |
159 |
| - if key and not isinstance(key, (keys.RSA2048, keys.RSA2048Public)): |
160 |
| - raise Exception("Encryption with sign only available with RSA") |
161 |
| - img.create(key, enckey) |
162 |
| - |
163 |
| - if pad: |
164 |
| - img.pad_to(slot_size) |
165 |
| - |
166 |
| - img.save(outfile) |
167 |
| - |
168 |
| - |
169 |
| -class AliasesGroup(click.Group): |
170 |
| - |
171 |
| - _aliases = { |
172 |
| - "create": "sign", |
173 |
| - } |
174 |
| - |
175 |
| - def list_commands(self, ctx): |
176 |
| - cmds = [k for k in self.commands] |
177 |
| - aliases = [k for k in self._aliases] |
178 |
| - return sorted(cmds + aliases) |
179 |
| - |
180 |
| - def get_command(self, ctx, cmd_name): |
181 |
| - rv = click.Group.get_command(self, ctx, cmd_name) |
182 |
| - if rv is not None: |
183 |
| - return rv |
184 |
| - if cmd_name in self._aliases: |
185 |
| - return click.Group.get_command(self, ctx, self._aliases[cmd_name]) |
186 |
| - return None |
187 |
| - |
188 |
| - |
189 |
| -@click.command(cls=AliasesGroup, |
190 |
| - context_settings=dict(help_option_names=['-h', '--help'])) |
191 |
| -def imgtool(): |
192 |
| - pass |
193 |
| - |
194 |
| - |
195 |
| -imgtool.add_command(keygen) |
196 |
| -imgtool.add_command(getpub) |
197 |
| -imgtool.add_command(sign) |
198 |
| - |
| 17 | +from imgtool import main |
199 | 18 |
|
200 | 19 | if __name__ == '__main__':
|
201 |
| - imgtool() |
| 20 | + main.imgtool() |
0 commit comments