imgtool: fixed keys/general.py to pass existing unittests
keys.KeyClass._emit is able to use 'file' parameter not as a file
but some object (not only sys.stdout but io.StringIO, like by
tests).
Fixed all explicit checks for sys.stdio usage in favor of
io.TextIOBase, also improve a single unit test to cover
also all the changed methods.
Signed-off-by: Denis Mingulov <denis@mingulov.com>
diff --git a/scripts/imgtool/keys/general.py b/scripts/imgtool/keys/general.py
index b415f47..4ff700c 100644
--- a/scripts/imgtool/keys/general.py
+++ b/scripts/imgtool/keys/general.py
@@ -2,22 +2,39 @@
# SPDX-License-Identifier: Apache-2.0
+import binascii
+import io
+import os
import sys
from cryptography.hazmat.primitives.hashes import Hash, SHA256
AUTOGEN_MESSAGE = "/* Autogenerated by imgtool.py, do not edit. */"
+class FileHandler(object):
+ def __init__(self, file, *args, **kwargs):
+ self.file_in = file
+ self.args = args
+ self.kwargs = kwargs
+
+ def __enter__(self):
+ if isinstance(self.file_in, (str, bytes, os.PathLike)):
+ self.file = open(self.file_in, *self.args, **self.kwargs)
+ else:
+ self.file = self.file_in
+ return self.file
+
+ def __exit__(self, *args):
+ if self.file != self.file_in:
+ self.file.close()
+
+
class KeyClass(object):
def _emit(self, header, trailer, encoded_bytes, indent, file=sys.stdout,
len_format=None):
- if file and file is not sys.stdout:
- with open(file, 'w') as file:
- self._emit_to_output(header, trailer, encoded_bytes, indent,
- file, len_format)
- else:
+ with FileHandler(file, 'w') as file:
self._emit_to_output(header, trailer, encoded_bytes, indent,
- sys.stdout, len_format)
+ file, len_format)
def _emit_to_output(self, header, trailer, encoded_bytes, indent, file,
len_format):
@@ -33,6 +50,16 @@
if len_format is not None:
print(len_format.format(len(encoded_bytes)), file=file)
+ def _emit_raw(self, encoded_bytes, file):
+ with FileHandler(file, 'wb') as file:
+ try:
+ # file.buffer is not part of the TextIOBase API
+ # and may not exist in some implementations.
+ file.buffer.write(encoded_bytes)
+ except AttributeError:
+ # raw binary data, can be for example io.BytesIO
+ file.write(encoded_bytes)
+
def emit_c_public(self, file=sys.stdout):
self._emit(
header="const unsigned char {}_pub_key[] = {{"
@@ -58,20 +85,12 @@
file=file)
def emit_raw_public(self, file=sys.stdout):
- if file and file is not sys.stdout:
- with open(file, 'wb') as file:
- file.write(self.get_public_bytes())
- else:
- sys.stdout.buffer.write(self.get_public_bytes())
+ self._emit_raw(self.get_public_bytes(), file=file)
def emit_raw_public_hash(self, file=sys.stdout):
digest = Hash(SHA256())
digest.update(self.get_public_bytes())
- if file and file is not sys.stdout:
- with open(file, 'wb') as file:
- file.write(digest.finalize())
- else:
- sys.stdout.buffer.write(digest.finalize())
+ self._emit_raw(digest.finalize(), file=file)
def emit_rust_public(self, file=sys.stdout):
self._emit(
@@ -83,11 +102,8 @@
file=file)
def emit_public_pem(self, file=sys.stdout):
- if file and file is not sys.stdout:
- with open(file, 'w') as file:
- print(str(self.get_public_pem(), 'utf-8'), file=file, end='')
- else:
- print(str(self.get_public_pem(), 'utf-8'), file=sys.stdout, end='')
+ with FileHandler(file, 'w') as file:
+ print(str(self.get_public_pem(), 'utf-8'), file=file, end='')
def emit_private(self, minimal, format, file=sys.stdout):
self._emit(
diff --git a/scripts/imgtool/keys/rsa_test.py b/scripts/imgtool/keys/rsa_test.py
index 4b2106a..7610106 100644
--- a/scripts/imgtool/keys/rsa_test.py
+++ b/scripts/imgtool/keys/rsa_test.py
@@ -63,15 +63,34 @@
for key_size in RSA_KEY_SIZES:
k = RSA.generate(key_size=key_size)
+ pubpem = io.StringIO()
+ k.emit_public_pem(pubpem)
+ self.assertIn("BEGIN PUBLIC KEY", pubpem.getvalue())
+ self.assertIn("END PUBLIC KEY", pubpem.getvalue())
+
ccode = io.StringIO()
k.emit_c_public(ccode)
self.assertIn("rsa_pub_key", ccode.getvalue())
self.assertIn("rsa_pub_key_len", ccode.getvalue())
+ hashccode = io.StringIO()
+ k.emit_c_public_hash(hashccode)
+ self.assertIn("rsa_pub_key_hash", hashccode.getvalue())
+ self.assertIn("rsa_pub_key_hash_len", hashccode.getvalue())
+
rustcode = io.StringIO()
k.emit_rust_public(rustcode)
self.assertIn("RSA_PUB_KEY", rustcode.getvalue())
+ # raw data - bytes
+ pubraw = io.BytesIO()
+ k.emit_raw_public(pubraw)
+ self.assertTrue(len(pubraw.getvalue()) > 0)
+
+ hashraw = io.BytesIO()
+ k.emit_raw_public_hash(hashraw)
+ self.assertTrue(len(hashraw.getvalue()) > 0)
+
def test_emit_pub(self):
"""Basic sanity check on the code emitters, from public key."""
pubname = self.tname("public.pem")