167 lines
3.8 KiB
Python
Executable File
167 lines
3.8 KiB
Python
Executable File
#
|
|
# Copyright (C) 2025 Amlogic, Inc. All rights reserved.
|
|
#
|
|
|
|
#!/usr/bin/python3 -B
|
|
|
|
import argparse
|
|
import binascii
|
|
import os
|
|
import pathlib
|
|
import struct
|
|
import sys
|
|
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives import serialization
|
|
from cryptography.hazmat.primitives.asymmetric import padding
|
|
from cryptography.hazmat.primitives.asymmetric import utils
|
|
|
|
#
|
|
# Variables
|
|
#
|
|
|
|
#
|
|
# Functions
|
|
#
|
|
|
|
def stderr_print(*arg, **kwargs):
|
|
print(*arg, file=sys.stderr, *kwargs);
|
|
return
|
|
|
|
#
|
|
# Classes
|
|
#
|
|
|
|
class DemoSign(object):
|
|
|
|
#
|
|
# Static variables
|
|
#
|
|
|
|
__FMT_HEADER__ = '<QIHH'
|
|
__MAGIC64__ = int.from_bytes(bytes.fromhex('58a848a4843ff07d'), byteorder='little')
|
|
|
|
#
|
|
# Initialization
|
|
#
|
|
|
|
def __init__(self, autorun=False):
|
|
self._args = None
|
|
|
|
if autorun:
|
|
self._parseArgs()
|
|
self._run()
|
|
return
|
|
|
|
#
|
|
# Methods (non-public)
|
|
#
|
|
|
|
def _parseArgs(self):
|
|
parser = argparse.ArgumentParser(description="Demo Sign")
|
|
parser.add_argument('INFILE_PRIVKEY', help='Input file of private RSA key')
|
|
|
|
self._args = parser.parse_args()
|
|
# print(self._args)
|
|
# print(vars(self._args))
|
|
# print(vars(self._args)['INFILE_PRIVKEY'])
|
|
# print(self._args.INFILE_PRIVKEY)
|
|
|
|
def _run(self):
|
|
self._recv_header()
|
|
self._recv_payload()
|
|
|
|
self._process()
|
|
|
|
self._send_header()
|
|
self._send_payload()
|
|
return
|
|
|
|
def _recv_header(self):
|
|
sz_hdr = struct.calcsize(type(self).__FMT_HEADER__)
|
|
buf = sys.stdin.buffer.read(sz_hdr)
|
|
assert len(buf) == sz_hdr
|
|
|
|
(
|
|
magic,
|
|
self._sz_payload,
|
|
self._seq_no,
|
|
self._flags
|
|
) = struct.unpack(type(self).__FMT_HEADER__, buf)
|
|
# stderr_print('magic: 0x{:x}'.format(magic))
|
|
assert magic == type(self).__MAGIC64__
|
|
assert (magic % 2) == 0
|
|
assert (self._seq_no % 2) == 0
|
|
|
|
return
|
|
|
|
def _recv_payload(self):
|
|
self._buf_digest = sys.stdin.buffer.read(self._sz_payload)
|
|
assert len(self._buf_digest) == self._sz_payload
|
|
return
|
|
|
|
def _process(self):
|
|
# Load private key
|
|
with open(self._args.INFILE_PRIVKEY, "rb") as key_file:
|
|
stderr_print('RSA key file: {}'.format(self._args.INFILE_PRIVKEY))
|
|
private_key = serialization.load_pem_private_key(
|
|
key_file.read(),
|
|
password=None,
|
|
)
|
|
del key_file
|
|
|
|
# Signing
|
|
assert len(self._buf_digest) == 32
|
|
# stderr_print(self._buf_digest.hex())
|
|
sig = private_key.sign(
|
|
self._buf_digest,
|
|
padding.PSS(
|
|
mgf=padding.MGF1(hashes.SHA256()),
|
|
salt_length=32
|
|
),
|
|
utils.Prehashed(hashes.SHA256())
|
|
)
|
|
self._buf_sig = sig
|
|
return
|
|
|
|
def _send_header(self):
|
|
sz_hdr = struct.calcsize(type(self).__FMT_HEADER__)
|
|
|
|
self._seq_no += 1
|
|
sz_payload = len(self._buf_sig)
|
|
buf = struct.pack(type(self).__FMT_HEADER__,
|
|
type(self).__MAGIC64__ + 1,
|
|
sz_payload,
|
|
self._seq_no,
|
|
0
|
|
)
|
|
assert len(buf) == sz_hdr
|
|
|
|
sz_out = sys.stdout.buffer.write(buf)
|
|
assert sz_out == sz_hdr
|
|
|
|
return
|
|
|
|
def _send_payload(self):
|
|
# stderr_print(self._buf_sig.hex())
|
|
sz_out = sys.stdout.buffer.write(self._buf_sig)
|
|
with open('test.sig', 'wb') as file2:
|
|
file2.write(self._buf_sig)
|
|
file2.close()
|
|
|
|
assert sz_out == len(self._buf_sig)
|
|
return
|
|
|
|
#
|
|
# Methods (public)
|
|
#
|
|
|
|
#
|
|
# Main
|
|
#
|
|
|
|
if __name__ == '__main__':
|
|
DemoSign(True)
|
|
|
|
# vim: set expandtab ai tabstop=4 shiftwidth=4 softtabstop=-1:
|