general.py

"""General key class."""

# SPDX-License-Identifier: Apache-2.0

import os
import sys

from cryptography.hazmat.primitives.hashes import SHA256, Hash

AUTOGEN_MESSAGE = "/* Autogenerated by imgtool.py, do not edit. */"


class FileHandler:
    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:
    def _emit(self, header, trailer, encoded_bytes, indent, file=sys.stdout,
              len_format=None):
        with FileHandler(file, 'w') as file:
            self._emit_to_output(header, trailer, encoded_bytes, indent,
                                     file, len_format)

    def _emit_to_output(self, header, trailer, encoded_bytes, indent, file,
                        len_format):
        print(AUTOGEN_MESSAGE, file=file)
        print(header, end='', file=file)
        for count, b in enumerate(encoded_bytes):
            if count % 8 == 0:
                print("\n" + indent, end='', file=file)
            else:
                print(" ", end='', file=file)
            print(f"0x{b:02x},", end='', file=file)
        print("\n" + trailer, file=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, name_suffix: str = ""):
        self._emit(
                header=f"const unsigned char {self.shortname()}_pub_key{name_suffix}[] = {{"
                       ,
                trailer="};",
                encoded_bytes=self.get_public_bytes(),
                indent="    ",
                len_format=f"const unsigned int {self.shortname()}_pub_key{name_suffix}_len = {{}};"
                           ,
                file=file)

    def emit_c_public_hash(self, file=sys.stdout, name_suffix: str = ""):
        digest = Hash(SHA256())
        digest.update(self.get_public_bytes())
        self._emit(
                header=f"const unsigned char {self.shortname()}_pub_key_hash{name_suffix}[] = {{"
                       ,
                trailer="};",
                encoded_bytes=digest.finalize(),
                indent="    ",
                len_format=("const unsigned int "
                            f"{self.shortname()}_pub_key_hash{name_suffix}_len = {{}};"),
                file=file)

    def emit_raw_public(self, file=sys.stdout):
        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())
        self._emit_raw(digest.finalize(), file=file)

    def emit_rust_public(self, file=sys.stdout, name_suffix: str = ""):
        self._emit(
                header=f"static {self.shortname().upper()}_PUB_KEY{name_suffix.upper()}: &[u8] = &["
                       ,
                trailer="];",
                encoded_bytes=self.get_public_bytes(),
                indent="    ",
                file=file)

    def emit_public_pem(self, file=sys.stdout):
        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(
                header="const unsigned char enc_priv_key[] = {",
                trailer="};",
                encoded_bytes=self.get_private_bytes(minimal, format),
                indent="    ",
                len_format="const unsigned int enc_priv_key_len = {};",
                file=file)