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 <[email protected]>

Author: mingulov

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 @@ def _emit_to_output(self, header, trailer, encoded_bytes, indent, file,
         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 @@ def emit_c_public_hash(self, file=sys.stdout):
                 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 @@ def emit_rust_public(self, file=sys.stdout):
                 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(

scripts/imgtool/keys/rsa_test.py

@@ -63,15 +63,34 @@ def test_emit(self):
         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")