import base64
import struct
import filecmp
from keyczar import util

# need pyasn for DER parsing and generating
from pyasn1.type import univ
from pyasn1.codec.der import decoder, encoder

def read_int(buffer, i):
    "Read 32bit integer from buffer."

    (l,) = struct.unpack('!I', buffer[i:i + 4])
    i += 4
    return (l, i)

def read_chunk(buffer, i):
    "Read chunk from buffer."

    # first grab length of chunk
    (l, i) = read_int(buffer, i)
    if l > 1000000:
        # just in case... if this happens, then something is way off
        raise Exception("got chunk length of %d, that's certainly too long" % l)

    # read chunk of length l
    #(s,) = struct.unpack('!' + '%ds' % l, buffer[i:i + l])
    (s,) = struct.unpack('!{0:d}s'.format(l), buffer[i:i + l])
    i += l
    return (s, i)

def read_rsa_pub(filename):
    """Read RSA public key file. Structure:

       ssh-rsa base64data user@host

       base64data: [7]ssh-rsa[len][e-data][len][n-data]
    """

    [prefix, data, host] = file(filename, 'r').read().split()
    raw = base64.b64decode(data)

    # read type string
    i = 0
    (s, i) = read_chunk(raw, i)
    if s != 'ssh-rsa':
        raise Exception("expected string 'ssh-rsa' but got '%s'" % s)

    # grab e
    (s, i) = read_chunk(raw, i)
    #e = unpack_bigint(s)
    e = util.BytesToLong(s)
    # grab n
    (s, i) = read_chunk(raw, i)
    #n = unpack_bigint(s)
    n = util.BytesToLong(s)
    return (n, e, host)

def write_rsa_pub(filename, n, e, host):
    """Write RSA public key file. Structure:

       ssh-rsa base64data user@host

       base64data: [7]ssh-rsa[len][e-data][len][n-data]
    """

    #e_str = pack_bigint(e)
    #n_str = pack_bigint(n)
    e_str = util.BigIntToBytes(e)
    n_str = util.BigIntToBytes(n)
    
    # pack e and n properly into the raw data
    raw = struct.pack('!I7sI{0:d}sI{1:d}s'.format(len(e_str), len(n_str)), 7, 'ssh-rsa',
                                        len(e_str), e_str, len(n_str), n_str)
    # assemble file content and save it
    content = "ssh-rsa {0!s} {1!s}\n".format(base64.b64encode(raw), host)
    file(filename, 'w').write(content)

def read_rsa_pri(filename):
    """Read RSA private key file. Structure:

       -----BEGIN RSA PRIVATE KEY-----
       base64data
       -----END RSA PRIVATE KEY-----

       base64data DER structure:

       RSAPrivateKey ::= SEQUENCE {
         version Version,
         modulus INTEGER,         -- n
         publicExponent INTEGER,  -- e
         privateExponent INTEGER, -- d
         prime1 INTEGER,          -- p
         prime2 INTEGER,          -- q
         exponent1 INTEGER,       -- d mod (p - 1)
         exponent2 INTEGER,       -- d mod (q - 1)
         coefficient INTEGER      -- q^-1 mod p
       }
    """

    # grab only the lines between the --- * --- lines, glue them together
    data = ''.join(filter(lambda x: x and x[0] != '-',
                          file(filename, 'r').read().split('\n')))
    # decode from base64
    raw = base64.b64decode(data)
    # parse DER structure
    der = decoder.decode(raw)
    (version, n, e, d, p, q, e1, e2, c) = (int(x) for x in der[0])

    return (n, e, d, p, q, e1, e2, c)

def write_rsa_pri(filename, n, e, d, p, q, e1, e2, c):
    """Write RSA private key file. Structure:

       -----BEGIN RSA PRIVATE KEY-----
       base64data
       -----END RSA PRIVATE KEY-----

       base64data DER structure:

       RSAPrivateKey ::= SEQUENCE {
         version Version,
         modulus INTEGER,         -- n
         publicExponent INTEGER,  -- e
         privateExponent INTEGER, -- d
         prime1 INTEGER,          -- p
         prime2 INTEGER,          -- q
         exponent1 INTEGER,       -- d mod (p - 1)
         exponent2 INTEGER,       -- d mod (q - 1)
         coefficient INTEGER      -- q^-1 mod p
       }
    """

    seq = (
            univ.Integer(0),
            univ.Integer(n),
            univ.Integer(e),
            univ.Integer(d),
            univ.Integer(p),
            univ.Integer(q),
            univ.Integer(e1),
            univ.Integer(e2),
            univ.Integer(c),
          )
    struct = univ.Sequence()
    for i in xrange(len(seq)):
        struct.setComponentByPosition(i, seq[i])

    # build DER structure
    raw = encoder.encode(struct)
    # encode to base64
    data = base64.b64encode(raw)

    # chop data up into lines of certain width
    width = 64
    chopped = [data[i:i + width] for i in xrange(0, len(data), width)]
    # assemble file content
    content = """-----BEGIN RSA PRIVATE KEY-----
{0}
-----END RSA PRIVATE KEY-----
""".format('\n'.join(chopped))
    file(filename, 'w').write(content)

if __name__ == '__main__' :
    ssh_keys_directory='/home/guillaume/.ssh/'
    print 'Testing public key reading...'
    (n,e,host)=read_rsa_pub(ssh_keys_directory + 'id_rsa.pub')
    write_rsa_pub(ssh_keys_directory + 'id_rsa_test.pub',n,e,host)
    if filecmp.cmp(ssh_keys_directory + 'id_rsa.pub',ssh_keys_directory + 'id_rsa_test.pub'):
        print 'test succesful'
        
    print 'Testing private key reading...'
    (n, e, d, p, q, e1, e2, c)=read_rsa_pri('/home/guillaume/.ssh/id_rsa')
    write_rsa_pri('/home/guillaume/.ssh/id_rsa_test',n, e, d, p, q, e1, e2, c)
    if filecmp.cmp(ssh_keys_directory + 'id_rsa',ssh_keys_directory + 'id_rsa_test'):
        print 'test succesful'