import argparse
import time
import sys

# Use pycryptodome's RSA library
# See https://pycryptodome.readthedocs.io/en/latest/src/public_key/rsa.html#Crypto.PublicKey.RSA.construct  # noqa: E501
from Cryptodome.PublicKey import RSA


KEY_ELEMENTS = ["PRIV_EXP", "PRIME_1", "PRIME_2", "PUB_EXP"]


def sparse_str_to_hex(s):
    return int(s.replace('"', "").replace(" ", ""), 16)


def main(infile=None):
    lines = []
    # Checks, if interactive session or reading from a file
    if sys.stdin.isatty():
        if infile.name != "<stdin>":
            # Reads specification file
            with open(infile.name, "r") as f:
                lines = f.readlines()
        else:
            exit("Supply data via STDIN or file!")
    else:
        lines = sys.stdin.readlines()

    construct_key(lines)


def construct_key(lines):
    keymaterial = {}

    # Populates keymaterial
    for line in lines:
        # Tokenizes line
        parts = line.split("=", 1)
        name = parts[0]
        val = parts[-1]

        # Stores line in dict
        if name in KEY_ELEMENTS:
            keymaterial[name] = sparse_str_to_hex(val)

    if "PRIME_1" in keymaterial and "PRIME_2" in keymaterial:
        keymaterial["MODULUS"] = keymaterial["PRIME_1"] * keymaterial["PRIME_2"]  # noqa: E501
    else:
        exit("Key specification is incomplete.")

    # Constructs RSA key by given params
    rsa_key = RSA.construct(
        (
            keymaterial["MODULUS"],  # n
            keymaterial["PUB_EXP"],  # e
            keymaterial["PRIV_EXP"],  # d
            keymaterial["PRIME_1"],  # p
            keymaterial["PRIME_2"],  # q
        )
    )

    # Assign it to a variable to keep it in memory
    exported_rsa_key = rsa_key.exportKey()  # noqa: F841

    while True:
        time.sleep(1)


def parse_args():
    parser = argparse.ArgumentParser(
            description="Construct a RSA key according to specified key material and sleep forever."  # noqa: E501
    )
    parser.add_argument(
        "keyspec", nargs="?", type=argparse.FileType("r"), default=sys.stdin
    )
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = parse_args()
    construct_key(args.keyspec)
