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(