#!/usr/bin/env python
# @ CommonUtility.py
# Common utility script
#
# Copyright (c) 2016 - 2021, Intel Corporation. All rights reserved.<BR>
# SPDX-License-Identifier: BSD-2-Clause-Patent
#
##

import os
import sys
import shutil
import subprocess
import string
from ctypes import ARRAY, c_char, c_uint16, c_uint32, \
        c_uint8, Structure, sizeof
from importlib.machinery import SourceFileLoader
from SingleSign import single_sign_gen_pub_key


# Key types  defined should match with cryptolib.h
PUB_KEY_TYPE = {
    "RSA": 1,
    "ECC": 2,
    "DSA": 3,
    }

# Signing type schemes  defined should match with cryptolib.h
SIGN_TYPE_SCHEME = {
    "RSA_PKCS1": 1,
    "RSA_PSS": 2,
    "ECC": 3,
    "DSA": 4,
    }

# Hash values defined should match with cryptolib.h
HASH_TYPE_VALUE = {
    "SHA2_256": 1,
    "SHA2_384": 2,
    "SHA2_512": 3,
    "SM3_256": 4,
    }

# Hash values defined should match with cryptolib.h
HASH_VAL_STRING = dict(map(reversed, HASH_TYPE_VALUE.items()))

AUTH_TYPE_HASH_VALUE = {
    "SHA2_256": 1,
    "SHA2_384": 2,
    "SHA2_512": 3,
    "SM3_256": 4,
    "RSA2048SHA256": 1,
    "RSA3072SHA384": 2,
    }

HASH_DIGEST_SIZE = {
    "SHA2_256": 32,
    "SHA2_384": 48,
    "SHA2_512": 64,
    "SM3_256": 32,
    }


class PUB_KEY_HDR (Structure):
    _pack_ = 1
    _fields_ = [
        ('Identifier', ARRAY(c_char, 4)),      # signature ('P', 'U', 'B', 'K')
        ('KeySize',    c_uint16),              # Length of Public Key
        ('KeyType',    c_uint8),               # RSA or ECC
        ('Reserved',   ARRAY(c_uint8, 1)),
        ('KeyData',    ARRAY(c_uint8, 0)),
        ]

    def __init__(self):
        self.Identifier = b'PUBK'


class SIGNATURE_HDR (Structure):
    _pack_ = 1
    _fields_ = [
        ('Identifier', ARRAY(c_char, 4)),
        ('SigSize',    c_uint16),
        ('SigType',    c_uint8),
        ('HashAlg',    c_uint8),
        ('Signature',  ARRAY(c_uint8, 0)),
        ]

    def __init__(self):
        self.Identifier = b'SIGN'


class LZ_HEADER(Structure):
    _pack_ = 1
    _fields_ = [
        ('signature',       ARRAY(c_char, 4)),
        ('compressed_len',  c_uint32),
        ('length',          c_uint32),
        ('version',         c_uint16),
        ('svn',             c_uint8),
        ('attribute',       c_uint8)
    ]
    _compress_alg = {
        b'LZDM': 'Dummy',
        b'LZ4 ': 'Lz4',
        b'LZMA': 'Lzma',
    }


def print_bytes(data, indent=0, offset=0, show_ascii=False):
    bytes_per_line = 16
    printable = ' ' + string.ascii_letters + string.digits + string.punctuation
    str_fmt = '{:s}{:04x}: {:%ds} {:s}' % (bytes_per_line * 3)
    bytes_per_line
    data_array = bytearray(data)
    for idx in range(0, len(data_array), bytes_per_line):
        hex_str = ' '.join(
            '%02X' % val for val in data_array[idx:idx + bytes_per_line])
        asc_str = ''.join('%c' % (val if (chr(val) in printable) else '.')
                          for val in data_array[idx:idx + bytes_per_line])
        print(str_fmt.format(
            indent * ' ',
            offset + idx, hex_str,
            ' ' + asc_str if show_ascii else ''))


def get_bits_from_bytes(bytes, start, length):
    if length == 0:
        return 0
    byte_start = (start) // 8
    byte_end = (start + length - 1) // 8
    bit_start = start & 7
    mask = (1 << length) - 1
    val = bytes_to_value(bytes[byte_start:byte_end + 1])
    val = (val >> bit_start) & mask
    return val


def set_bits_to_bytes(bytes, start, length, bvalue):
    if length == 0:
        return
    byte_start = (start) // 8
    byte_end = (start + length - 1) // 8
    bit_start = start & 7
    mask = (1 << length) - 1
    val = bytes_to_value(bytes[byte_start:byte_end + 1])
    val &= ~(mask << bit_start)
    val |= ((bvalue & mask) << bit_start)
    bytes[byte_start:byte_end+1] = value_to_bytearray(
        val,
        byte_end + 1 - byte_start)


def value_to_bytes(value, length):
    return value.to_bytes(length, 'little')


def bytes_to_value(bytes):
    return int.from_bytes(bytes, 'little')


def value_to_bytearray(value, length):
    return bytearray(value_to_bytes(value, length))

# def value_to_bytearray (value, length):
    return bytearray(value_to_bytes(value, length))


def get_aligned_value(value, alignment=4):
    if alignment != (1 << (alignment.bit_length() - 1)):
        raise Exception(
            'Alignment (0x%x) should to be power of 2 !' % alignment)
    value = (value + (alignment - 1)) & ~(alignment - 1)
    return value


def get_padding_length(data_len, alignment=4):
    new_data_len = get_aligned_value(data_len, alignment)
    return new_data_len - data_len


def get_file_data(file, mode='rb'):
    return open(file, mode).read()


def gen_file_from_object(file, object):
    open(file, 'wb').write(object)


def gen_file_with_size(file, size):
    open(file, 'wb').write(b'\xFF' * size)


def check_files_exist(base_name_list, dir='', ext=''):
    for each in base_name_list:
        if not os.path.exists(os.path.join(dir, each + ext)):
            return False
    return True


def load_source(name, filepath):
    mod = SourceFileLoader(name, filepath).load_module()
    return mod


def get_openssl_path():
    if os.name == 'nt':
        if 'OPENSSL_PATH' not in os.environ:
            openssl_dir = "C:\\Openssl\\bin\\"
            if os.path.exists(openssl_dir):
                os.environ['OPENSSL_PATH'] = openssl_dir
            else:
                os.environ['OPENSSL_PATH'] = "C:\\Openssl\\"
                if 'OPENSSL_CONF' not in os.environ:
                    openssl_cfg = "C:\\Openssl\\openssl.cfg"
                    if os.path.exists(openssl_cfg):
                        os.environ['OPENSSL_CONF'] = openssl_cfg
        openssl = os.path.join(
            os.environ.get('OPENSSL_PATH', ''),
            'openssl.exe')
    else:
        # Get openssl path for Linux cases
        openssl = shutil.which('openssl')

    return openssl


def run_process(arg_list, print_cmd=False, capture_out=False):
    sys.stdout.flush()
    if os.name == 'nt' and os.path.splitext(arg_list[0])[1] == '' and \
       os.path.exists(arg_list[0] + '.exe'):
        arg_list[0] += '.exe'
    if print_cmd:
        print(' '.join(arg_list))

    exc = None
    result = 0
    output = ''
    try:
        if capture_out:
            output = subprocess.check_output(arg_list).decode()
        else:
            result = subprocess.call(arg_list)
    except Exception as ex:
        result = 1
        exc = ex

    if result:
        if not print_cmd:
            print('Error in running process:\n  %s' % ' '.join(arg_list))
        if exc is None:
            sys.exit(1)
        else:
            raise exc

    return output


# Adjust hash type algorithm based on Public key file
def adjust_hash_type(pub_key_file):
    key_type = get_key_type(pub_key_file)
    if key_type == 'RSA2048':
        hash_type = 'SHA2_256'
    elif key_type == 'RSA3072':
        hash_type = 'SHA2_384'
    else:
        hash_type = None

    return hash_type


def rsa_sign_file(
      priv_key, pub_key, hash_type, sign_scheme,
      in_file, out_file, inc_dat=False, inc_key=False):

    bins = bytearray()
    if inc_dat:
        bins.extend(get_file_data(in_file))


# def single_sign_file(priv_key, hash_type, sign_scheme, in_file, out_file):

    out_data = get_file_data(out_file)

    sign = SIGNATURE_HDR()
    sign.SigSize = len(out_data)
    sign.SigType = SIGN_TYPE_SCHEME[sign_scheme]
    sign.HashAlg = HASH_TYPE_VALUE[hash_type]

    bins.extend(bytearray(sign) + out_data)
    if inc_key:
        key = gen_pub_key(priv_key, pub_key)
        bins.extend(key)

    if len(bins) != len(out_data):
        gen_file_from_object(out_file, bins)


def get_key_type(in_key):

    # Check in_key is file or key Id
    if not os.path.exists(in_key):
        key = bytearray(gen_pub_key(in_key))
    else:
        # Check for public key in binary format.
        key = bytearray(get_file_data(in_key))

    pub_key_hdr = PUB_KEY_HDR.from_buffer(key)
    if pub_key_hdr.Identifier != b'PUBK':
        pub_key = gen_pub_key(in_key)
        pub_key_hdr = PUB_KEY_HDR.from_buffer(pub_key)

    key_type = next(
        (key for key,
            value in PUB_KEY_TYPE.items() if value == pub_key_hdr.KeyType))
    return '%s%d' % (key_type, (pub_key_hdr.KeySize - 4) * 8)


def get_auth_hash_type(key_type, sign_scheme):
    if key_type == "RSA2048" and sign_scheme == "RSA_PKCS1":
        hash_type = 'SHA2_256'
        auth_type = 'RSA2048_PKCS1_SHA2_256'
    elif key_type == "RSA3072" and sign_scheme == "RSA_PKCS1":
        hash_type = 'SHA2_384'
        auth_type = 'RSA3072_PKCS1_SHA2_384'
    elif key_type == "RSA2048" and sign_scheme == "RSA_PSS":
        hash_type = 'SHA2_256'
        auth_type = 'RSA2048_PSS_SHA2_256'
    elif key_type == "RSA3072" and sign_scheme == "RSA_PSS":
        hash_type = 'SHA2_384'
        auth_type = 'RSA3072_PSS_SHA2_384'
    else:
        hash_type = ''
        auth_type = ''
    return auth_type, hash_type


# def single_sign_gen_pub_key(in_key, pub_key_file=None):


def gen_pub_key(in_key, pub_key=None):

    keydata = single_sign_gen_pub_key(in_key, pub_key)

    publickey = PUB_KEY_HDR()
    publickey.KeySize = len(keydata)
    publickey.KeyType = PUB_KEY_TYPE['RSA']

    key = bytearray(publickey) + keydata

    if pub_key:
        gen_file_from_object(pub_key, key)

    return key


def decompress(in_file, out_file, tool_dir=''):
    if not os.path.isfile(in_file):
        raise Exception("Invalid input file '%s' !" % in_file)

    # Remove the Lz Header
    fi = open(in_file, 'rb')
    di = bytearray(fi.read())
    fi.close()

    lz_hdr = LZ_HEADER.from_buffer(di)
    offset = sizeof(lz_hdr)
    if lz_hdr.signature == b"LZDM" or lz_hdr.compressed_len == 0:
        fo = open(out_file, 'wb')
        fo.write(di[offset:offset + lz_hdr.compressed_len])
        fo.close()
        return

    temp = os.path.splitext(out_file)[0] + '.tmp'
    if lz_hdr.signature == b"LZMA":
        alg = "Lzma"
    elif lz_hdr.signature == b"LZ4 ":
        alg = "Lz4"
    else:
        raise Exception("Unsupported compression '%s' !" % lz_hdr.signature)

    fo = open(temp, 'wb')
    fo.write(di[offset:offset + lz_hdr.compressed_len])
    fo.close()

    compress_tool = "%sCompress" % alg
    if alg == "Lz4":
        try:
            cmdline = [
                os.path.join(tool_dir, compress_tool),
                "-d",
                "-o", out_file,
                temp]
            run_process(cmdline, False, True)
        except Exception:
            msg_string = "Could not find/use CompressLz4 tool, " \
                        "trying with python lz4..."
            print(msg_string)
            try:
                import lz4.block
                if lz4.VERSION != '3.1.1':
                    msg_string = "Recommended lz4 module version " \
                                "is '3.1.1'," + lz4.VERSION \
                                + " is currently installed."
                    print(msg_string)
            except ImportError:
                msg_string = "Could not import lz4, use " \
                            "'python -m pip install lz4==3.1.1' " \
                            "to install it."
                print(msg_string)
                exit(1)
            decompress_data = lz4.block.decompress(get_file_data(temp))
            with open(out_file, "wb") as lz4bin:
                lz4bin.write(decompress_data)
    else:
        cmdline = [
            os.path.join(tool_dir, compress_tool),
            "-d",
            "-o", out_file,
            temp]
        run_process(cmdline, False, True)
    os.remove(temp)


def compress(in_file, alg, svn=0, out_path='', tool_dir=''):
    if not os.path.isfile(in_file):
        raise Exception("Invalid input file '%s' !" % in_file)

    basename, ext = os.path.splitext(os.path.basename(in_file))
    if out_path:
        if os.path.isdir(out_path):
            out_file = os.path.join(out_path, basename + '.lz')
        else:
            out_file = os.path.join(out_path)
    else:
        out_file = os.path.splitext(in_file)[0] + '.lz'

    if alg == "Lzma":
        sig = "LZMA"
    elif alg == "Tiano":
        sig = "LZUF"
    elif alg == "Lz4":
        sig = "LZ4 "
    elif alg == "Dummy":
        sig = "LZDM"
    else:
        raise Exception("Unsupported compression '%s' !" % alg)

    in_len = os.path.getsize(in_file)
    if in_len > 0:
        compress_tool = "%sCompress" % alg
        if sig == "LZDM":
            shutil.copy(in_file, out_file)
            compress_data = get_file_data(out_file)
        elif sig == "LZ4 ":
            try:
                cmdline = [
                    os.path.join(tool_dir, compress_tool),
                    "-e",
                    "-o", out_file,
                    in_file]
                run_process(cmdline, False, True)
                compress_data = get_file_data(out_file)
            except Exception:
                msg_string = "Could not find/use CompressLz4 tool, " \
                            "trying with python lz4..."
                print(msg_string)
                try:
                    import lz4.block
                    if lz4.VERSION != '3.1.1':
                        msg_string = "Recommended lz4 module version " \
                                    "is '3.1.1', " + lz4.VERSION \
                                    + " is currently installed."
                        print(msg_string)
                except ImportError:
                    msg_string = "Could not import lz4, use " \
                                "'python -m pip install lz4==3.1.1' " \
                                "to install it."
                    print(msg_string)
                    exit(1)
                compress_data = lz4.block.compress(
                    get_file_data(in_file),
                    mode='high_compression')
        elif sig == "LZMA":
            cmdline = [
                os.path.join(tool_dir, compress_tool),
                "-e",
                "-o", out_file,
                in_file]
            run_process(cmdline, False, True)
            compress_data = get_file_data(out_file)
    else:
        compress_data = bytearray()

    lz_hdr = LZ_HEADER()
    lz_hdr.signature = sig.encode()
    lz_hdr.svn = svn
    lz_hdr.compressed_len = len(compress_data)
    lz_hdr.length = os.path.getsize(in_file)
    data = bytearray()
    data.extend(lz_hdr)
    data.extend(compress_data)
    gen_file_from_object(out_file, data)

    return out_file