quantcrypt/prepare_submission.py at main · codespree/quantcrypt · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
# Run this script to automatically zip the correct certificates for submission
# Note: The certficates shall be generated by running the rust testcases as follows:
# Part 1: Generate non-ipd certificates
# cargo test --release
# Part 1: Generate ipd certificates
# cargo test --release --features ipd

import os
import zipfile
import shutil

"""
    Prepare submission zip for R4 or R5 certificates (DER only).
    Usage:
        prepare_dsa_submission()          # R4
        prepare_dsa_submission(r5=True)   # R5
"""
def prepare_dsa_submission(cert_format=".der", r5=True):
    import os, shutil, zipfile

    version = "r5" if r5 else "r4"
    print(f"Preparing {version.upper()} certificate submission zip file")

    cert_path = f"./artifacts/{version}_certs"
    assert os.path.exists(cert_path), f"Certificates not generated, please run Rust test cases first"

    submission_dir = "./artifacts/submission"
    os.makedirs(submission_dir, exist_ok=True)

    artifacts_certs = os.path.join(submission_dir, "artifacts_certs")
    os.makedirs(artifacts_certs, exist_ok=True)

    # Copy DER certificates
    for root, _, files in os.walk(cert_path):
        for file in files:
            if file.endswith(cert_format):
                shutil.copy(os.path.join(root, file), artifacts_certs)

    # Zip
    zip_name = f"artifacts_certs_{version}.zip"
    zip_path = os.path.join(submission_dir, zip_name)
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for root, _, files in os.walk(artifacts_certs):
            for file in files:
                full_path = os.path.join(root, file)
                zipf.write(full_path, os.path.relpath(full_path, artifacts_certs))

    shutil.rmtree(artifacts_certs)
    print(f"Submission zip file created at {zip_path}")

def prepare_cms_submission(cert_format: str = ".der"):

    print("Preparing cms submission zip file")

    cert_path = f"./artifacts/v2_cms"
    assert os.path.exists(cert_path), "Certificates not generated, please run Rust test cases first"

    submission_dir = "./artifacts/submission"
    os.makedirs(submission_dir, exist_ok=True)

    # Create temporary folder for submission certificates
    artifacts_certs = os.path.join(submission_dir, "artifacts_certs_cms")
    os.makedirs(artifacts_certs, exist_ok=True)

    # Copy the DER certificates in cert_path to artifacts_certs
    for root, dirs, files in os.walk(cert_path):
        for file in files:
            if file.endswith(cert_format):
                shutil.copy(os.path.join(root, file), artifacts_certs)


    # Zip all files in artifacts_certs
    zip_file = os.path.join(submission_dir, "artifacts_cms_v2.zip")
    with zipfile.ZipFile(zip_file, 'w') as zipf:
        for root, dirs, files in os.walk(artifacts_certs):
            for file in files:
                zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), artifacts_certs))

    # Remove the temporary folder after zipping
    shutil.rmtree(artifacts_certs)

    print(f"Submission zip file created at {zip_file}")

if __name__ == "__main__":
    prepare_dsa_submission()
    prepare_cms_submission()