Skip to content

Commit b5115c4

Browse files
tomato42tomato42
committed
scripts to generate IETF standard key formats and use them
1 parent e7e7a42 commit b5115c4

File tree

3 files changed

+304
-0
lines changed

3 files changed

+304
-0
lines changed

‎scripts/decaps.py‎

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

‎scripts/encaps.py‎

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

‎scripts/keygen.py‎

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

0 commit comments

Comments
 (0)