Skip to content

Commit c30a719

Browse files
tomato42tomato42
committed
scripts to generate IETF standard key formats and use them
1 parent 970a760 commit c30a719

File tree

3 files changed

+296
-0
lines changed

3 files changed

+296
-0
lines changed

‎scripts/decaps.py‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
import sys
2+
import getopt
3+
from kyber_py.ml_kem.pkcs import (
4+
dk_from_pem,
5+
dk_from_der,
6+
)
7+
8+
9+
def help_msg():
10+
print(f"""Usage: {sys.argv[0]} [options]
11+
--dk FILE Decapsulation key file name
12+
--dk-form FORM Decapsulation key format name: "DER" or "PEM, with "PEM"
13+
being the default.
14+
--secret FILE Name of file to write the derived secret to
15+
--ciphertext FILE Name of file to read the ciphertext from
16+
--help This message
17+
""")
18+
19+
20+
def main():
21+
dk_file = None
22+
dk_form = "PEM"
23+
secret_file = None
24+
ciphertext_file = None
25+
26+
argv = sys.argv[1:]
27+
opts, args = getopt.getopt(
28+
argv,
29+
"",
30+
["dk=", "dk-form=", "secret=", "ciphertext=", "help"],
31+
)
32+
33+
for opt, arg in opts:
34+
if opt == "--dk":
35+
dk_file = arg
36+
elif opt == "--dk-form":
37+
dk_form = arg
38+
elif opt == "--secret":
39+
secret_file = arg
40+
elif opt == "--ciphertext":
41+
ciphertext_file = arg
42+
elif opt == "--help":
43+
help_msg()
44+
sys.exit(0)
45+
else:
46+
print(f"Unrecognised option: {opt}")
47+
sys.exit(1)
48+
49+
if args:
50+
print(f"Unrecognised options: {args}")
51+
52+
if not dk_file:
53+
print("Specify encapsulation key name with --dk option")
54+
help_msg()
55+
sys.exit(1)
56+
57+
if not secret_file:
58+
print("Specify file name to write the derived secret to")
59+
help_msg()
60+
sys.exit(1)
61+
62+
if not ciphertext_file:
63+
print("Specify file name to read the ciphertext from")
64+
help_msg()
65+
sys.exit(1)
66+
67+
if dk_form == "PEM":
68+
dk_dec_func = dk_from_pem
69+
elif dk_form == "DER":
70+
dk_dec_func = dk_from_der
71+
else:
72+
print(f"Wrong name of key format: {dk_form}")
73+
help_msg()
74+
sys.exit(1)
75+
76+
with open(dk_file, "rb") as file:
77+
kem, dk, _, _ = dk_dec_func(file.read())
78+
79+
with open(ciphertext_file, "rb") as file:
80+
ciphertext = file.read()
81+
82+
secret = kem.decaps(dk, ciphertext)
83+
84+
with open(secret_file, "wb") as file:
85+
file.write(secret)
86+
87+
88+
if __name__ == "__main__":
89+
main()

‎scripts/encaps.py‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
import sys
2+
import getopt
3+
from kyber_py.ml_kem.pkcs import (
4+
ek_from_pem,
5+
ek_from_der,
6+
)
7+
8+
def help_msg():
9+
print(f"""Usage: {sys.argv[0]} [options]
10+
--ek FILE Encapsulation key file name
11+
--ek-form FORM Encapsulation key format name: "DER" or "PEM, with "PEM"
12+
being the default.
13+
--secret FILE Name of file to write the derived secret to
14+
--ciphertext FILE Name of file to write the ciphertext to
15+
--help This message
16+
""")
17+
18+
19+
def main():
20+
ek_file = None
21+
ek_form = "PEM"
22+
secret_file = None
23+
ciphertext_file = None
24+
25+
argv = sys.argv[1:]
26+
opts, args = getopt.getopt(
27+
argv,
28+
"",
29+
["ek=", "ek-form=", "secret=", "ciphertext=", "help"],
30+
)
31+
32+
for opt, arg in opts:
33+
if opt == "--ek":
34+
ek_file = arg
35+
elif opt == "--ek-form":
36+
ek_form = arg
37+
elif opt == "--secret":
38+
secret_file = arg
39+
elif opt == "--ciphertext":
40+
ciphertext_file = arg
41+
elif opt == "--help":
42+
help_msg()
43+
sys.exit(0)
44+
else:
45+
print(f"Unrecognised option: {opt}")
46+
sys.exit(1)
47+
48+
if args:
49+
print(f"Unrecognised options: {args}")
50+
51+
if not ek_file:
52+
print("Specify encapsulation key name with --ek option")
53+
help_msg()
54+
sys.exit(1)
55+
56+
if not secret_file:
57+
print("Specify file name to write the derived secret to")
58+
help_msg()
59+
sys.exit(1)
60+
61+
if not ciphertext_file:
62+
print("Specify file name to write the ciphertext to")
63+
help_msg()
64+
sys.exit(1)
65+
66+
if ek_form == "PEM":
67+
ek_dec_func = ek_from_pem
68+
elif ek_form == "DER":
69+
ek_dec_func = ek_from_der
70+
else:
71+
print(f"Wrong name of key format: {ek_form}")
72+
help_msg()
73+
sys.exit(1)
74+
75+
with open(ek_file, "rb") as file:
76+
kem, ek = ek_dec_func(file.read())
77+
78+
secret, ciphertext = kem.encaps(ek)
79+
80+
with open(secret_file, "wb") as file:
81+
file.write(secret)
82+
83+
with open(ciphertext_file, "wb") as file:
84+
file.write(ciphertext)
85+
86+
87+
if __name__ == "__main__":
88+
main()

‎scripts/keygen.py‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
import sys
2+
import getopt
3+
import os
4+
from kyber_py.ml_kem import ML_KEM_512, ML_KEM_768, ML_KEM_1024
5+
6+
from kyber_py.ml_kem.pkcs import (
7+
ek_to_der,
8+
ek_to_pem,
9+
dk_to_der,
10+
dk_to_pem,
11+
)
12+
13+
14+
def help_msg():
15+
print(f"""Usage: {sys.argv[0]} [options]
16+
17+
--dk FILE Decapsulation key file name
18+
--dk-form FORM Decapsulation key format name: "DER" or "PEM", with "PEM"
19+
being the default.
20+
--dk-cont CONT Decapsulation key contents: "seed", "expanded", or "both",
21+
with "both" being the default.
22+
--ek FILE Encapsulation key file name, none by default
23+
--ek-form FORM Encapsulation key format name: "DER" or "PEM", with "PEM"
24+
being the default.
25+
--kem NAME Name of the KEM to use: ML-KEM-512, ML-KEM-768, or ML-KEM-1024
26+
--help This message
27+
""")
28+
29+
def main():
30+
ek_file = None
31+
ek_form = "PEM"
32+
dk_file = None
33+
dk_form = "PEM"
34+
dk_cont = None
35+
kem = None
36+
37+
argv = sys.argv[1:]
38+
opts, args = getopt.getopt(
39+
argv,
40+
"",
41+
["dk=", "dk-form=", "dk-cont=", "ek=", "ek-form=", "kem=", "help"]
42+
)
43+
for opt, arg in opts:
44+
if opt == "--dk":
45+
dk_file = arg
46+
elif opt == "--dk-form":
47+
dk_form = arg
48+
elif opt == "--dk-cont":
49+
dk_cont = arg
50+
elif opt == "--ek":
51+
ek_file = arg
52+
elif opt == "--ek-form":
53+
ek_form = arg
54+
elif opt == "--kem":
55+
kem = arg
56+
elif opt == "--help":
57+
help_msg()
58+
sys.exit(0)
59+
else:
60+
print(f"unrecognised option: {opt}")
61+
sys.exit(1)
62+
63+
if args:
64+
print(f"unrecognised options: {args}")
65+
sys.exit(1)
66+
67+
if not dk_file:
68+
print("Specify output file with --dk option")
69+
help_msg()
70+
sys.exit(1)
71+
72+
if not kem:
73+
print("Specify the kem to generate with --kem option")
74+
help_msg()
75+
sys.exit(1)
76+
77+
if kem == "ML-KEM-512":
78+
kem = ML_KEM_512
79+
elif kem == "ML-KEM-768":
80+
kem = ML_KEM_768
81+
elif kem == "ML-KEM-1024":
82+
kem = ML_KEM_1024
83+
else:
84+
print(f"Unrecognised KEM name: {kem}")
85+
help_msg()
86+
sys.exit(1)
87+
88+
if ek_form == "PEM":
89+
ek_out_func = ek_to_pem
90+
elif ek_form == "DER":
91+
ek_out_func = ek_to_der
92+
else:
93+
print(f"Unrecognised ek format: {ek_form}")
94+
help_msg()
95+
sys.exit(1)
96+
97+
if dk_form == "PEM":
98+
dk_out_func = dk_to_pem
99+
elif dk_form == "DER":
100+
dk_out_func = dk_to_der
101+
else:
102+
print(f"Unrecognised dk format: {dk_form}")
103+
help_msg()
104+
sys.exit(1)
105+
106+
seed = os.urandom(64)
107+
108+
ek, dk = kem.key_derive(seed)
109+
110+
with open(dk_file, "wb") as file:
111+
file.write(dk_out_func(kem, dk, seed, dk_cont))
112+
113+
if ek_file:
114+
with open(ek_file, "wb") as file:
115+
file.write(ek_out_func(kem, ek))
116+
117+
118+
if __name__ == "__main__":
119+
main()

0 commit comments

Comments
 (0)