#!/usr/bin/env python3
# coding: utf-8
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
import os
import shutil
import zipfile
import tempfile
import logging
import argparse
import subprocess
import shlex

from flask import Flask, request, send_file

app = Flask(__name__)
logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s  %(levelname)s:%(message)s")
certs_dir = None
sign_pwd = None


@app.route('/sign', methods=['POST'])
def cms_sign():
    temp_dir = tempfile.TemporaryDirectory()
    os.chdir(temp_dir.name)
    cert_files = ["signer.pem.enc", "ts_signer.pem.enc", "tsa.cnf", "rootca.crl", "rootca.pem"]
    for file in cert_files:
        shutil.copy2(os.path.join(certs_dir, file), temp_dir.name)

    files = request.files.getlist("file")
    form_data = request.form
    metadata = form_data.get("metadata")
    if not metadata or not files:
        logging.error("No metadata or file")

    filename = None
    for file in files:
        filename = file.filename
        file_path = os.path.join(temp_dir.name, filename)
        file.save(file_path)
        break

    logging.info("===Request Reviced=====")


    signer = shutil.which("hpm_signer")
    sign_cmd = (f"{signer} -s signer.pem.enc -t ts_signer.pem.enc -T tsa.cnf "
            f" -i {file_path} -o {filename}.cms -p {sign_pwd}")
    run_cmd = shlex.split(sign_cmd)
    subprocess.run(run_cmd, check=True)

    verifier = shutil.which("hpm_verify")
    verify_cmd = f"{verifier} -r rootca.pem -C cms.crl.pem -s {filename}.cms -c {file_path}"
    run_cmd = shlex.split(verify_cmd)
    subprocess.run(run_cmd, check=True)
    logging.info("===Sign Succeeded=====")

    output_zip = os.path.join(temp_dir.name, "signed.zip")
    with zipfile.ZipFile(output_zip, 'w') as zipf:
        for file_name in ["signed.cms", "rootca.crl"]:
            zipf.write(os.path.join(temp_dir.name, file_name), file_name)
    return send_file(output_zip)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Start cms sign server", 
        add_help=True, formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("-certs", "--certs_dir", required=True, help="certs directory")
    parser.add_argument("-sp", "--sign_pwd", required=True, help="encryption password")
    args = parser.parse_args()
    sign_pwd = args.sign_pwd
    certs_dir = os.path.realpath(os.path.join(os.getcwd(), args.certs_dir))
    logging.info("Start cms sign server.............")
    app.run(debug=True, host="0.0.0.0", port=8098)
