<reponame>d-hrytsenko/virgil-iotkit<gh_stars>0
import json
import os
from collections import namedtuple

from waiting import wait, TimeoutExpired
from django.conf import settings

from sandbox.libs.cmd_process import CmdProcess
from sandbox.libs.logger import logger_by_name
from sandbox.libs.utils import find_files


class FWSigner:
    _signer_output = namedtuple('FWSignerOutput', 'prog_file update_file')

    def __init__(self,
                 fw_source: str,
                 version: str,
                 manufacturer: str,
                 model: str):
        self.fw_source = fw_source
        self.version = version
        self.manufacturer = manufacturer
        self.device_model = model
        self.cfg_path = _get_signer_config()

    def sign(self):
        # Calculate output file size
        file_size = os.stat(self.fw_source).st_size + 50000  # add some extra space

        # Run signer
        cmd = (
            f'virgil-firmware-signer '
            f'--input "{self.fw_source}" '
            f'--config "{self.cfg_path}" '
            f'--file-size {file_size} '
            f'--fw-version {self.version} '
            f'--manufacturer "{self.manufacturer}" '
            f'--model {self.device_model} '
            f'--chunk-size 64000'
        )
        process = CmdProcess(cmd,
                             log_path=os.path.join(settings.SANDBOX_LOGS_DIR, 'virgil-firmware-signer.log'))
        process.run()
        try:
            wait(lambda: process.return_code == 0, timeout_seconds=10)
        except TimeoutExpired:
            raise Exception(f'Failed to sign {self.fw_source} firmware.\n'
                            f'Return code of Virgil Firmware Signer: {process.return_code}')

        return self._signer_output(prog_file=f'{self.fw_source}_Prog.bin',
                                   update_file=f'{self.fw_source}_Update.bin')


def _get_signer_config():
    config_path = os.path.join(settings.SANDBOX_TEMP_DIR, 'signer-config.json')
    # since keys are generated once - return config if exists
    if os.path.exists(config_path):
        return config_path

    auth_private_key, *_ = find_files(settings.SANDBOX_TP_DIR, 'auth_.*key', regex=True)
    fw_private_key, *_ = find_files(settings.SANDBOX_TP_DIR, 'firmware_.*key', regex=True)

    data = [
        {'path': auth_private_key, 'key_type': 1},
        {'path': fw_private_key, 'key_type': 3},
    ]

    with open(config_path, 'w') as f:
        json.dump(data, f, ensure_ascii=False)

    return config_path
