import os
import zipfile
import shutil
def extract_correct_certs(version):
assert version in ["r3", "r4"], "Unknown submission version"
cert_format = "pem" if version == "r3" else "der"
cert_path = f"./artifacts/{version}_certs"
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)
artifacts_certs = os.path.join(submission_dir, f"artifacts_{version}_certs")
os.makedirs(artifacts_certs, exist_ok=True)
ipd_certs_path = os.path.join(cert_path, "ipd")
non_ipd_certs_path = os.path.join(cert_path, "non-ipd")
for cert in os.listdir(ipd_certs_path):
if cert.endswith(f".{cert_format}"):
shutil.copy(os.path.join(ipd_certs_path, cert), artifacts_certs)
for cert in os.listdir(non_ipd_certs_path):
cert_in_r3 = os.path.join(artifacts_certs, cert)
if not os.path.exists(cert_in_r3) and cert.endswith(f".{cert_format}"):
shutil.copy(os.path.join(non_ipd_certs_path, cert), artifacts_certs)
zip_filename = os.path.join(submission_dir, f"artifacts_{version}_certs.zip")
with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zipf:
for root, dirs, files in os.walk(artifacts_certs):
for file in files:
full_path = os.path.join(root, file)
relative_path = os.path.relpath(full_path, artifacts_certs)
zipf.write(full_path, relative_path)
shutil.rmtree(artifacts_certs)
if __name__ == "__main__":
extract_correct_certs("r3")
extract_correct_certs("r4")