# RADICAL SEMICONDUCTOR presents
#            _                    _ _ 
#  _ __ ___ (_)_  ____      _____| | |
# | '_ ` _ \| \ \/ /\ \ /\ / / _ \ | |
# | | | | | | |>  <  \ V  V /  __/ | |
# |_| |_| |_|_/_/\_\  \_/\_/ \___|_|_|
# 
#       a post-quantum challenge
#
# 
# NOTE: pubkey.p and cipher.p, output by this program, are available at:
#
#    https://radicalsemiconductor.com/resources/pubkey.p
#    https://radicalsemiconductor.com/resources/cipher.p
#
# Go to those links to download the key and cipher output by this file.
#
# DESCRIPTION:
# I did some research on lattice crypto methods this weekend. Not sure what
# all the hubbub with new methods is about, one method I found from a few
# years back seems to work great! I added a little twist to make sure the
# keys are good, but it shouldn't be necessary anyway.
#
# RULES:
# Find the FLAG! The below code was used to generate pubkey.p and cipher.p.
# The format of the flag is `radical{...}`, where `...` consists of letters,
# numbers, and underscores.
#
# CONTACT:
# If you find something wrong, or have feedback, etc., please reach out to
# info@radicalsemiconductor.com and we'll try to get back to you as soon
# as we can. 
#
# Additionally, we're looking to bring on a crypto advisor! If you believe
# you or someone you know would be a good fit, please reach out to the
# same address.
# 
#
# Once you find the flag, email it to info@radicalsemiconductor.com. The first
# winner will receive a Radical Semiconductor t-shirt and a $50 Amazon gift card!
# 
# UPDATE: first prize has been taken (faster than we thought!). We'll be tracking
# first 25 solves, so if you get the flag, email us at the above address and
# let us know if you'd like to be recognized on our website.
#
# Enjoy!
#
# --The Radical Team


import pickle

FLAG = b'radical{????????????????????????}'


"""Generate keys."""
def ggh_keygen(n):
    # while we're picking a key
    while True:
        # make our private key
        c = (4 * round(sqrt(n) + 1))
        rand = matrix(ZZ, n, n, [choice(range(-4, 4)) for _ in range(n**2)])
        R = c * matrix.identity(ZZ, n) + rand

        # make our public key
        B = copy(R)
        for i in range(2 * n):
            T = matrix.identity(n)
            col = vector([choice([-1, 1] + [0] * 5) for _ in range(n)])
            col[i % n] = 1
            T[:, i % n] = col
            B *= T

        # check that our key is "safe" mod 2 and 3            
        C = B.change_ring(Integers(2))
        D = B.change_ring(Integers(3))
        cdim = C.right_kernel().dimension()
        ddim = D.right_kernel().dimension()

        # if so, we're done
        if cdim >= 1 and ddim >= 1:
            break

    return B, R


"""Encrypt a message."""
def ggh_encrypt(m, B):
    # zero pad
    n = B.nrows()
    m += b'\0' * (n - len(m))

    # encode 
    x = vector(ZZ, [a - 128 for a in list(m)])

    # get error
    e = vector(ZZ, [choice([-3, 3]) for _ in range(n)])

    # encrypt
    return B * x + e


"""Decrypt a message."""
def ggh_decrypt(c, B, R):
    # perform Babai rounding
    v = B.inverse() * R * (R.inverse() * c).apply_map(round)

    # decode
    m = bytes([a + 128 for a in v if a + 128 != 0])
    return m


m = FLAG
B, R = ggh_keygen(150)
c = ggh_encrypt(m, B)

with open('pubkey.p', 'wb') as f:
    pickle.dump(B, f)

with open('cipher.p', 'wb') as f:
    pickle.dump(c, f)