import bson, scrypt

from Crypto.Cipher import AES
from Crypto.Util import Counter
from Crypto.PublicKey import RSA
from Crypto.PublicKey import DSA
from salsa20 import Salsa20_keystream
import os, struct, time, hashlib, hashlib, random


class EncryptedScreenshot:
    def __init__(self, metadata,id=None,signer=None,password_encryptor=None):
        def rand(len):
            return ''.join(
                random.choice("1234567890ABCDEFGHIJKLMNOPQRSTUWVXYZabcdefghijklmnopqrstuwvxyz") for _ in range(len))

        self.password = rand(16)
        print "Passphrase %s" % str(bytearray(self.passphrase())).encode("hex")
	self.id = id
	if id is None:
        	self.id = rand(8)
        self.metadata_encryption = True
        self.metadata = metadata
        self.signer = signer
        self.password_encryptor = password_encryptor

    def caesar_encrypted_password(self):
        caesar_key = reduce(lambda sum,charcode: sum+charcode, map(ord,list(self.id)))
        caesar_plaintext = map(ord, self.password)
        caesar_ciphertext = Caesar().encrypt(caesar_plaintext,caesar_key)
        return ''.join(map(chr,caesar_ciphertext))

    def passphrase(self):
        # new ScryptParameters(64, 8, 1,32, new Uint8List.fromList(new List<int>()))
        print "Password units: %s" % (map(ord, self.password.encode("utf-8")),)
        sha = hashlib.sha256()
        sha.update(self.password)
        return sha.digest()  # scrypt.hash(self.password.encode("utf-8"), '', 64, 8, 1, 32)

    def assemble(self, file):
        (image_digest, encrypted_image) = self.encrypt(file)
        self.metadata["hash"] = image_digest
        unencrypted_metadata = bson.dumps(self.metadata)
        if len(unencrypted_metadata) % 16 != 0:
            unencrypted_metadata += ' ' * (16 - len(unencrypted_metadata) % 16)
        (encryptor, iv) = self.encryptor(len(unencrypted_metadata))
        encrypted_metadata = []

        encrypted_metadata += encryptor(unencrypted_metadata)

        encrypted_metadata = iv + str(bytearray(encrypted_metadata))

        print "Metadata: %s" % str(encrypted_metadata).encode("base64").replace("\n", "")
        print "%s %s" % (str(encrypted_metadata[:16]).encode("hex"), str(encrypted_metadata[16:]).encode("hex"))
        print "Unencrypted: %s" % (unencrypted_metadata.encode("hex"))
        print "Password %s" % self.password

        print bson.loads(unencrypted_metadata)

        fields = {
            "image": encrypted_image,
            "metadata_encryption": self.metadata_encryption,
            "metadata": encrypted_metadata if self.metadata_encryption else self.metadata
        }

        if self.signer is not None:
            fields["signature"] = self.signer.signature(encrypted_metadata if self.metadata_encryption else bson.dumps(self.metadata))

        if self.password_encryptor is not None:
            fields["password"] = self.password_encryptor.encrypt_password(self.password)

        return bson.dumps(fields)

    def encryptor(self,length=0):
        iv = os.urandom(16)
        ctr = Counter.new(128, initial_value=long(iv.encode("hex"), 16))
        print "IV: %s" % iv.encode("hex")
        cipher = AES.new(self.passphrase(), AES.MODE_CTR, counter=ctr)

        #salsa
        # keysstream = bytearray(Salsa20_keystream(length,iv[:8],self.passphrase()))
        # offset={"offset": 0}
        def encrypt(data):
            #Salsa20
            # array = bytearray(data)
            # for i in range(len(array)):
            #     array[i] ^= keysstream[i+offset["offset"]]
            # offset["offset"] += len(array)
            # return str(array)
            return cipher.encrypt(data)

        return (encrypt, iv)

    def encrypt(self, file):
        filesize = os.path.getsize(file)
        (encryptor, iv) = self.encryptor(filesize)
        binary = []
        digest = hashlib.sha256()
        with open(file, 'rb') as infile:
            # binary += struct.pack('<Q', filesize)
            while True:
                chunk = infile.read(64 * 1024)
                if len(chunk) == 0:
                    break
                elif len(chunk) % 16 != 0:
                    pass  # chunk += ' ' * (16 - len(chunk) % 16)
                digest.update(chunk)
                binary += encryptor(chunk)
        return (digest.digest(), iv + str(bytearray(binary)))

from Crypto.Signature import PKCS1_v1_5
from Crypto.Hash import SHA256

class Signer:

    @staticmethod
    def default():
        return Signer(RSA.importKey("-----BEGIN RSA PRIVATE KEY-----\nMIIJKAIBAAKCAgEAuTEHVWjtwFydG/VmHKr15qymIkwZIaNnQeKmBzzn+n87Zjdg\nkEDtCEYSq3W1VxlgUw1Dh8KTTr3K1LQwcDfx8vdArk99OX4mxqpYWRbdfEQk1wUD\n7ctr2E5eKQcCcbAl/OKYFa/PkY3KsBuns06Rlf85sHSZ7CfirqH6FoefrXDZVAJO\nmrQhTh/0qGlBy1f13c2a12sFNOmncLs8Vp4ftJxLr9GoE4vsQvX/6oQvi7PqWYmv\nlwmUUvhCPS/JSF8kCuvmvfgcF8OXOj1Q6itOLZDt8/ZuZeok7nS8JxhR73HzHR4a\nHWnqC8UkIx1CR/VvlZhp/7gbS/y1EORRaVlO3YTkF2UMTtF2BoCsY5glpealDV6F\nh9YZwq68APb5/cleuXu8tebYpo5yuwH9G8m0bp7mIrVUnaBPCE/Wsj6HFULniUi8\n0N8149cM2kLmm7c+JNC4QUupc3iy3qylqh1KFlU0SWPAinDbODSHqnPiWV0IHjC9\nmVeAmMDjg/6R7uaAubLz0rJso/odk0kfWYO98D43Hv7C+TG/If9oE7JdsgSCCfeT\nuaKE5XH3C3GegciDOyr6Sd/08KprQwvy7nZpJJxzpyr2MoJ/n8O6O4IqDt76EAZz\n76nUvJzPXSoUGy7942QOwywpRYB6NbrVpPUemUCsUttUJzhgHU70ElqE2hsCAwEA\nAQKCAgBVyMZS5BYt4zjLuLMiWHcr8ZHLlWCW2nB8bpbmB2qV+NX1FXXfnJOwH3WK\ngy4LOLAMCLrdal5+PSB1wzok4WVsXOhTVnL1bOBdao8eRZ83bwF5sNMPFgnQ2wx3\nCM4zB5uxNu2oLhaCoCHhzg2y25RI3T9F491hLWCt8E4+nXh0CTBC4Aa1VveOjvLw\nCEH4BbjYk+NEIkj5J7ruV20t2nVpXnUY7o9i7o/0vKtdYPn9RswpycPSvRc+D83v\nNfGRWrEBdiEeoo2w3GNzEPAoaQL/zIDNX7vn7bxJ0T1PBs79kdVCCa2jLmrCJQIG\nWglX9/NGYqHHWOscGeVt+7VSjZIf8+q32khXcheBUP++wsDC1YoxiT2aDwV7zX8d\nBhB6MMi7QUqkkwnuvXdyUIUy4J7lJHDsC5l4EYDmbhViTax8Qp42MobJ32tHoS1c\nhfN5w4Usqx7yGRL0vUQmbGrh7JbyayD6CFYiZOinhN0ma7dMqglQ4ch3HPeMct9T\nLskF+gSlol2gTCMYjl9Zwb2rG8uSthjXFd2CDCHbiqzzRhskHG2vWuhxvASP4yeD\nNbm85H8ZKqhejcWL4XlJIZK2ZbWUjdMDPCDDd3ipwQ/GUmSo6IaNxEFMCFLsk/xl\n2BGg99xDuacldKNUtf+aMv4HglW6at+TasVsmWXsF/cWR1JusQKCAQEA88JECQTE\nm8AIOaHY5rOWxkYCsVzVw9zRHOfIXbvMMr9k3OTE/9gNy35QTee5UYVEnjRuWP/r\n2hRJhbTfwP5LmcvsxN3cE3SkrlNCulo0TWD4O78W+PnEkSzN+1h39QOfA/TTN9pV\nnqh7Bf5u9kBioooPeUe6yU/eMlVa+mzwNfDC9ogS7o1/1Zy4ui4beo0XpIv5KNtV\nU5K2uEcx0jnAzhcd+IKzDNm7HuU82KCsYNUCdHIJOpoq0WqPTqa350LPWHqbZA9n\nOa+5Yo1ADMn/rMFohErWq2zjUlO3qFjyGBaVAQCYCeprmk56vk4gP44X2YkW524G\nHg/CWYITF42W7QKCAQEAwn3UDzS1LeCZq82TSlFPWO3dojh4fMwssQWzJkrkJg4N\nNlAGGteh7K2eIFjACk7bLe037F2RhPxk+RZEIy4Adr5oAcNLKXMIcaAx9Xca+JrX\nRjLAGsBuoOJFYBVm+/ZlwoVkhOnX1VOJ4Zt4TNOYcImrhZyeVjVfJI7ijTU61v5M\n2+6Ekc2pq6WI7/2hmM+UOw6xvroNvCx0neWVgAVdCWfWWmCkUqjsS8FJ2z9q/8or\nnFPKeAT/kcDg2N+8LiPeVfAqdl6GZ3SY3lZ+Q+JEJMAZEsFsbh0x5OT4750Y/MBz\nqQF1oU2HkV5xah5tAqECp37/FSIxPG2e/qklGqbMJwKCAQA/qkSUym8DcvNzAgeh\nsdQ2cwCvDF00q87eKteVS8B4nK6eVw3w+KAZvyMARf6MAS+71lvf2O/j5JqTVnt+\nb6yVPcLYZE3fsv50HYD3JSo175bGHBHG841YkYg3FpM2Jjz5Do/ALcFhEAYZuauH\nm6dnZoHGHmr1u2uAX/TEJYHTxM9XDpFLRx/VTfqdff9lUcaUDo8EEc1GZ7JbjRTB\nIJ5v9W+5gYF8YL3TwvHSbpy3KctbWJF6gdufGKLXWwn7V+LJ7DP7vV4rne7gg53v\n+S5TE6zVqKpK2cTBsuXsDVcl3XgLa/lJ1bVG9ypt2//mOZmXEowt4/mrDlingia0\nZyG9AoIBAAvBQSTeIakxLviER0exnGMalU+trBKW4IieeJJy6lEmfkPvGD9dEwtT\n/v+jfY/h7jtbYAroKNhpzMyFPP8HU73Zij25CPzrV/JOJVbZqQanjSVQAe5b/RO6\nOUkyYfQOk13rZATogUQCdY3uMJosG8WJfJVpvuQk+6gg/q5cLj86RXjxdI+/Z2bZ\nbWpb4P10ZlgfbjzBz/y0jyPP8XY4XMQpw3nY53WbghSeDRG0RopJNNOQ+ObCHpQ3\nnR9efmgngNOe3WIkGml8UZT2k1EVBttVxO4uRJhvKUnVUPpie1UjaJopXWuvXOc7\nVw1740jKZHetJEupsmIdmgN4Q6bVt/cCggEBAILm8bfc2OGbZgN4RT4315EkWekb\nMngMFETDDgn6p2jmkVjXDGku1oX502X7BKteZ+gD4Oqy8BtrvshyzvzoBerthrDj\nJ4M3t7uFC14mvRfdelRcKyed13SGzZNJxxxL7vNq2Me81yzBaMPiAZpDIN4awEzV\nPQwWMmogo6sEp2ND2NhvbP4Hxl8WXoI+3hH5hetE31kvBzDlWek4Wj059BeeIhjF\nBNUxIwbqlSwt+2VuY17LBzCe4Tnp3Ns+ptD9Fn2vm+Azl9qBY2VvqOcdp0R8tmeo\nwpjrWKwP9LkwCt9ZxVzCQBiIFYtpZ8a2EahqTbOlFEsvcDbTfQlf/lykhjA=\n-----END RSA PRIVATE KEY-----\n"),"PUBLIC")

    def __init__(self,privateKey,privateKeyId, mode=RSA):
        self.mode = mode
        self.privateKeyId = privateKeyId
        self._signer = PKCS1_v1_5.new(privateKey)

    def sign(self, data):
        digest = SHA256.new()
        digest.update(data)
        return self._signer.sign(digest)

    def signature(self,data):
        signed = self.sign(data)
        return {"signed-hash": signed,
                "signature-algorithm": "SHA-256/%s" % unicode("RSA" if self.mode == RSA else "DSA"),
                "key-id": unicode(self.privateKeyId)
                }


class EncryptedPassword:

    def __init__(self,publicKey,publicKeyId,mode=RSA):
        assert mode == RSA
        self.mode = mode
        self.publicKey = publicKey
        self.publicKeyId = publicKeyId

    def encrypt(self,password):
        return self.publicKey.encrypt(password.encode("utf-8"), os.urandom(32))[0]

    def encrypt_password(self,password):
        return {"algorithm": unicode( "RSA" if self.mode == RSA else "DSA"),
                "key-id": unicode(self.publicKeyId),
                "password": self.encrypt(password)
                }

class Caesar:

    def __init__(self,charset = map(ord,list("1234567890ABCDEFGHIJKLMNOPQRSTUWVXYZabcdefghijklmnopqrstuwvxyz"))):
        self.charset = charset

    def encrypt(self,plaintext,key):
         if key >= len(self.charset) or key < 0:
            key = key % len(self.charset)

         def addKey(charcode):
            pos = self.charset.index(charcode)
            if pos + key >= len(self.charset):
                return self.charset[pos+key - len(self.charset)]
            else:
                return self.charset[pos+key]
         return map(addKey,plaintext)

    def decrypt(self,ciphertext,key):
        return self.encrypt(ciphertext,key*-1)


if __name__ == "__main__":
    privateKey = " "
    publicKey = "LS0tLS1CRUdJTiBQVUJMSUMgS0VZLS0tLS0KTUlJQ0lqQU5CZ2txaGtpRzl3MEJBUUVGQUFPQ0FnOEFNSUlDQ2dLQ0FnRUF1VEVIVldqdHdGeWRHL1ZtSEtyMQo1cXltSWt3WklhTm5RZUttQnp6bituODdaamRna0VEdENFWVNxM1cxVnhsZ1V3MURoOEtUVHIzSzFMUXdjRGZ4Cjh2ZEFyazk5T1g0bXhxcFlXUmJkZkVRazF3VUQ3Y3RyMkU1ZUtRY0NjYkFsL09LWUZhL1BrWTNLc0J1bnMwNlIKbGY4NXNIU1o3Q2ZpcnFINkZvZWZyWERaVkFKT21yUWhUaC8wcUdsQnkxZjEzYzJhMTJzRk5PbW5jTHM4VnA0Zgp0SnhMcjlHb0U0dnNRdlgvNm9Rdmk3UHFXWW12bHdtVVV2aENQUy9KU0Y4a0N1dm12ZmdjRjhPWE9qMVE2aXRPCkxaRHQ4L1p1WmVvazduUzhKeGhSNzNIekhSNGFIV25xQzhVa0l4MUNSL1Z2bFpocC83Z2JTL3kxRU9SUmFWbE8KM1lUa0YyVU1UdEYyQm9Dc1k1Z2xwZWFsRFY2Rmg5WVp3cTY4QVBiNS9jbGV1WHU4dGViWXBvNXl1d0g5RzhtMApicDdtSXJWVW5hQlBDRS9Xc2o2SEZVTG5pVWk4ME44MTQ5Y00ya0xtbTdjK0pOQzRRVXVwYzNpeTNxeWxxaDFLCkZsVTBTV1BBaW5EYk9EU0hxblBpV1YwSUhqQzltVmVBbU1EamcvNlI3dWFBdWJMejBySnNvL29kazBrZldZTzkKOEQ0M0h2N0MrVEcvSWY5b0U3SmRzZ1NDQ2ZlVHVhS0U1WEgzQzNHZWdjaURPeXI2U2QvMDhLcHJRd3Z5N25acApKSnh6cHlyMk1vSi9uOE82TzRJcUR0NzZFQVp6NzZuVXZKelBYU29VR3k3OTQyUU93eXdwUllCNk5iclZwUFVlCm1VQ3NVdHRVSnpoZ0hVNzBFbHFFMmhzQ0F3RUFBUT09Ci0tLS0tRU5EIFBVQkxJQyBLRVktLS0tLQo="
    encrypted = EncryptedScreenshot({"owner": u"test", "format": u"jpg", "title": u"5x5 Red", "timestamp": int(time.time()*1000)},
                                    signer=Signer.default(),
                                    password_encryptor=EncryptedPassword(RSA.importKey(publicKey.decode("base64")),"PUBLIC"))
    fixed_id = encrypted.id #"W9u9Zm0u"
    print len(encrypted.passphrase())
    out = open("/home/marvin/Dokumente/IdeaProjects/EncryptedScreencloud/Frontend/WebApp/web/data/" + fixed_id, 'wb')
    assembled = encrypted.assemble("5x5red.png")
    out.write(assembled)
    b = bson.loads(assembled)
    del b["image"]
    print b
    #print assembled.encode("base64").replace("\n","")
    print "http://localhost:8080/#%s%s" % (fixed_id, encrypted.password)
    caesar_key = reduce(lambda sum,charcode: sum+charcode, map(ord,list(fixed_id)))
    caesar_plaintext = map(ord,encrypted.password)
    caesar_ciphertext = Caesar().encrypt(caesar_plaintext,caesar_key)
    assert caesar_plaintext == Caesar().decrypt(caesar_ciphertext,caesar_key)
    print "caesar_key=%s %s -> %s" % (caesar_key,encrypted.password,''.join(map(chr,caesar_ciphertext)))
    print "http://localhost:8080/#%s%sC" % (fixed_id, encrypted.caesar_encrypted_password())
    print "http://localhost:8080/#%s" % (fixed_id, )
    out.close()
