
from Crypto.Hash import SHA256
from Crypto.PublicKey import RSA


class Cipher():
    def __init__(self, keyfile, key_length=2048, e=65537, gen_key=True):
        """init Cipher object.
        """
        if gen_key==False:
            with open(keyfile, 'r') as f:
                self.rsa = RSA.importKey(f.read())
        else:
            self.rsa = RSA.generate(key_length, e=65537)
            with open(keyfile, 'w') as f:
                f.write(self.rsa.exportKey('PEM'))

    def encrypt(self, plain):
        """encrypt string into list of chuncks. maximum length of each chunk is
        250 bytes.
        """
        if self.rsa.can_encrypt():
            ciphertext = []
            for i in range(0, len(plain), 250):
                # The second argument is useless.
                ciphertext.append(self.rsa.encrypt(plain[i:i+250], 0)[0])

            return ciphertext
        else:
            print '[Error] can\'t encrypt'
            return None

    def decrypt(self, cipher):
        """input list of chuncks of cipher text.
        """
        plain = ''
        for c in cipher:
            plain += self.rsa.decrypt(c)

        return plain

    def get_fingerprint(self):
        """returns the base64 of fingerprint of this object, generated by SHA256.
        """
        sha = SHA256.new()
        sha.update(self.rsa.publickey().exportKey('DER'))
        return sha.digest()

    def get_publickey(self):
        """returns the public of this object.
        """
        return self.rsa.publickey().exportKey('DER')


def main():
    c = Cipher('mykey.pem', gen_key=True)
    cip = c.encrypt("""
            This is a good project, very very very fun,
            trying to get this message over 256 bytes to show that this object can
            handle message over 250 bytes.
            Who can help me think of more words to type ~
            I need more message.
            aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
            bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
            Hope this is enough.
            """
        )
    print 'ciphertext: ', cip
    print 'decrypted message: ' + c.decrypt(cip)


if __name__ == '__main__':
    main()
