import base64
from Crypto.Cipher import AES
import hashlib

class Crypto:
    def __init__(self, secret):
        self.secret = add_pad(secret[0:16])
        self.cipher = AES.new(self.secret, AES.MODE_ECB)

    def encrypt(self, msg):

        return base64.urlsafe_b64encode(
            self.cipher.encrypt(
                add_pad(
                    msg
                )
            )
        ).decode("utf-8")


    def decrypt(self, enc):

        return remove_pad(
            self.cipher.decrypt(
                base64.urlsafe_b64decode(
                    enc
                )
            ).decode("utf-8")
        )


def add_pad(string):
    """ Add spaces until length is multiple of 16 """
    while len(string)%16 != 0:
        string+=" "
    return string


def get_direct_token(share, filename):
    if not 'pass_hash' in share:
        return None
    return password_hash(
        share['pass_hash'] + filename
    )


def password_hash(string):
    if type(string) == str:
        string = string.encode("utf-8")
    return hashlib.sha1(
        string
    ).hexdigest()


def remove_pad(string):
    """ Remove spaces from right """
    return string.rstrip(" ")