#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-or-later
#
# Script that generates lib/crypto/fips-aes.h and lib/crypto/fips-sha.h
#
# Requires that python-cryptography be installed.
#
# Copyright 2025 Google LLC

import cryptography.hazmat.primitives.ciphers
import cryptography.hazmat.primitives.ciphers.aead
import cryptography.hazmat.primitives.cmac
import hashlib
import hmac


def print_static_u8_array_definition(file, name, value):
    print("", file=file)
    print(f"static const u8 {name}[] __initconst __maybe_unused = {{", file=file)
    for i in range(0, len(value), 8):
        line = "\t" + "".join(f"0x{b:02x}, " for b in value[i : i + 8])
        print(f"{line.rstrip()}", file=file)
    print("};", file=file)


def print_header(file):
    print("/* SPDX-License-Identifier: GPL-2.0-or-later */", file=file)
    print("/* This file was generated by: gen-fips-testvecs.py */", file=file)
    print("/* clang-format off */", file=file)
    print("", file=file)
    print("#include <linux/fips.h>", file=file)


def gen_aes_test_data(file):
    fips_test_data = b"fips test data\0\0"
    fips_test_ad = b"fips test ad\0\0\0\0"
    fips_test_iv = b"fips test iv\0\0\0\0"
    fips_test_key = b"fips test key\0\0\0"
    fips_test_xts_key = b"key1" + (b"\0" * 12) + b"key2" + (b"\0" * 12)

    print_header(file)
    print_static_u8_array_definition(file, "fips_test_data", fips_test_data)
    print_static_u8_array_definition(file, "fips_test_ad", fips_test_ad)
    print_static_u8_array_definition(file, "fips_test_iv", fips_test_iv)
    print_static_u8_array_definition(file, "fips_test_key", fips_test_key)
    print_static_u8_array_definition(file, "fips_test_xts_key", fips_test_xts_key)

    aes = cryptography.hazmat.primitives.ciphers.algorithms.AES(fips_test_key)

    # AES-CMAC
    aes_cmac = cryptography.hazmat.primitives.cmac.CMAC(aes)
    aes_cmac.update(fips_test_data)
    print_static_u8_array_definition(
        file, "fips_test_aes_cmac_value", aes_cmac.finalize()
    )

    # AES-ECB
    cipher = cryptography.hazmat.primitives.ciphers.Cipher(
        aes, cryptography.hazmat.primitives.ciphers.modes.ECB()
    )
    encryptor = cipher.encryptor()
    ctext = encryptor.update(fips_test_data) + encryptor.finalize()
    print_static_u8_array_definition(file, "fips_test_aes_ecb_ctext", ctext)

    # AES-CBC
    cipher = cryptography.hazmat.primitives.ciphers.Cipher(
        aes, cryptography.hazmat.primitives.ciphers.modes.CBC(fips_test_iv)
    )
    encryptor = cipher.encryptor()
    ctext = encryptor.update(fips_test_data) + encryptor.finalize()
    print_static_u8_array_definition(file, "fips_test_aes_cbc_ctext", ctext)

    # AES-CBC-CTS
    cipher = cryptography.hazmat.primitives.ciphers.Cipher(
        aes, cryptography.hazmat.primitives.ciphers.modes.CBC(fips_test_iv)
    )
    encryptor = cipher.encryptor()
    ctext = encryptor.update(fips_test_data * 2) + encryptor.finalize()
    ctext = ctext[16:32] + ctext[0:16]
    print_static_u8_array_definition(file, "fips_test_aes_cbc_cts_ctext", ctext)

    # AES-CTR
    cipher = cryptography.hazmat.primitives.ciphers.Cipher(
        aes, cryptography.hazmat.primitives.ciphers.modes.CTR(fips_test_iv)
    )
    encryptor = cipher.encryptor()
    ctext = encryptor.update(fips_test_data) + encryptor.finalize()
    print_static_u8_array_definition(file, "fips_test_aes_ctr_ctext", ctext)

    # AES-XTS
    cipher = cryptography.hazmat.primitives.ciphers.Cipher(
        cryptography.hazmat.primitives.ciphers.algorithms.AES(fips_test_xts_key),
        cryptography.hazmat.primitives.ciphers.modes.XTS(fips_test_iv),
    )
    encryptor = cipher.encryptor()
    ctext = encryptor.update(fips_test_data) + encryptor.finalize()
    print_static_u8_array_definition(file, "fips_test_aes_xts_ctext", ctext)

    # AES-GCM
    cipher = cryptography.hazmat.primitives.ciphers.aead.AESGCM(fips_test_key)
    ct_and_tag = cipher.encrypt(
        nonce=fips_test_iv[:12], data=fips_test_data, associated_data=fips_test_ad
    )
    print_static_u8_array_definition(
        file, "fips_test_aes_gcm_ctext_and_tag", ct_and_tag
    )

    # AES-CCM
    cipher = cryptography.hazmat.primitives.ciphers.aead.AESCCM(
        fips_test_key, tag_length=16
    )
    ct_and_tag = cipher.encrypt(
        nonce=fips_test_iv[:13], data=fips_test_data, associated_data=fips_test_ad
    )
    print_static_u8_array_definition(
        file, "fips_test_aes_ccm_ctext_and_tag", ct_and_tag
    )


def gen_sha_test_data(file):
    fips_test_data = b"fips test data\0\0"
    fips_test_key = b"fips test key\0\0\0"

    print_header(file)
    print_static_u8_array_definition(file, "fips_test_data", fips_test_data)
    print_static_u8_array_definition(file, "fips_test_key", fips_test_key)

    for alg in "sha1", "sha256", "sha512":
        ctx = hmac.new(fips_test_key, digestmod=alg)
        ctx.update(fips_test_data)
        print_static_u8_array_definition(
            file, f"fips_test_hmac_{alg}_value", ctx.digest()
        )

    print_static_u8_array_definition(
        file, "fips_test_sha3_256_value", hashlib.sha3_256(fips_test_data).digest()
    )


filename = "lib/crypto/fips-aes.h"
with open(filename, "w") as file:
    print(f"Generating {filename}")
    gen_aes_test_data(file)

filename = "lib/crypto/fips-sha.h"
with open(filename, "w") as file:
    print(f"Generating {filename}")
    gen_sha_test_data(file)