#!/usr/bin/env python3

import sys

assert sys.version_info[:2] >= (3,0), "This is Python 3 code"

def generate():
    import hashlib

    print("""\
# DO NOT EDIT DIRECTLY! Autogenerated by agenttestgen.py
#
# To regenerate, run
#     python3 agenttestgen.py > agenttestdata.py
#
# agenttestgen.py depends on the testcrypt system, so you must also
# have built testcrypt in the parent directory, or else set
# PUTTY_TESTCRYPT to point at a working implementation of it.

""")

    from testcrypt import (rsa_generate, dsa_generate, ecdsa_generate,
                           eddsa_generate, random_clear, random_queue,
                           ssh_key_public_blob, ssh_key_openssh_blob,
                           ssh_key_sign, rsa1_generate, rsa_ssh1_encrypt,
                           rsa_ssh1_public_blob, rsa_ssh1_private_blob_agent,
                           mp_from_bytes_be)
    from agenttest import (Key2, TestSig2, test_message_to_sign,
                           Key1, test_session_id)
    import ssh

    keygen2 = [
        ('RSA-1024', lambda: rsa_generate(1024, False),
         (ssh.SSH_AGENT_RSA_SHA2_256, ssh.SSH_AGENT_RSA_SHA2_512)),
        ('DSA-1024', lambda: dsa_generate(1024)),
        ('ECDSA-p256', lambda: ecdsa_generate(256)),
        ('Ed25519', lambda: eddsa_generate(256)),
    ]

    keys2 = []

    for record in keygen2:
        if len(record) == 2:
            record += ((),)
        comment, genfn, flaglist = record
        flaglist = (0,) + flaglist

        random_clear()
        random_queue(b''.join(hashlib.sha512('{}{:d}'.format(comment, j)
                                             .encode('ASCII')).digest()
                              for j in range(1000)))
        key = genfn()
        sigs = [TestSig2(flags, ssh_key_sign(key, test_message_to_sign, flags))
                for flags in flaglist]

        keys2.append(Key2(comment.encode("ASCII"),
                          ssh_key_public_blob(key),
                          sigs,
                          ssh_key_openssh_blob(key)))

    print("def key2examples(Key2, TestSig2):\n    return {!r}".format(keys2))

    keygen1 = [
        ('RSA-1024a', 1024),
        ('RSA-1024b', 1024),
        ('RSA-768c', 768),
        ('RSA-768d', 768),
    ]

    keys1 = []

    for comment, bits in keygen1:
        random_clear()
        random_queue(b''.join(hashlib.sha512('{}{:d}'.format(comment, j)
                                             .encode('ASCII')).digest()
                              for j in range(1000)))
        key = rsa1_generate(bits)
        preimage = b'Test128BitRSA1ChallengeCleartext'
        assert len(preimage) == 32
        challenge_bytes = rsa_ssh1_encrypt(preimage, key)
        assert len(challenge_bytes) > 0
        challenge = int(mp_from_bytes_be(challenge_bytes))
        response = hashlib.md5(preimage + test_session_id).digest()

        keys1.append(Key1(comment.encode("ASCII"),
                          rsa_ssh1_public_blob(key, 'exponent_first'),
                          challenge, response,
                          rsa_ssh1_private_blob_agent(key)))

    print("def key1examples(Key1):\n    return {!r}".format(keys1))

if __name__ == "__main__":
    generate()