# -*- coding: utf-8 -*-

import re
import random
from array import array
from os import remove
from base64 import standard_b64decode

from Crypto.Cipher import AES
from Crypto.Util import Counter

from module.common.json_layer import json
from module.plugins.Hoster import Hoster

#def getInfo(urls):
#    pass

class MegaNz(Hoster):
    __name__ = "MegaNz"
    __type__ = "hoster"
    __pattern__ = r"https?://([a-z0-9]+\.)?mega\.co\.nz/#!([a-zA-Z0-9!_\-]+)"
    __version__ = "0.1"
    __description__ = """mega.co.nz hoster plugin"""
    __author_name__ = ("RaNaN", )
    __author_mail__ = ("ranan@pyload.org", )

    API_URL = "https://g.api.mega.co.nz/cs?id=%d"
    FILE_SUFFIX = ".crypted"

    def b64_decode(self, data):
        return standard_b64decode(data.replace("-", "+").replace("_", "/")+ "=")

    def getCipherKey(self, key):
        """ Construct the cipher key from the given data """
        a = array("I", key)
        key_array = array("I", [a[0] ^ a[4], a[1] ^a[5], a[2] ^ a[6], a[3] ^a[7]])
        return key_array

    def callApi(self, **kwargs):
        """ Dispatch a call to the api, see https://mega.co.nz/#developers """
        # generate a session id, no idea where to obtain elsewhere
        uid = random.randint(10 << 9, 10 ** 10)

        resp = self.load(self.API_URL % uid, post=json.dumps([kwargs]))
        self.logDebug("Api Response: " + resp)
        return json.loads(resp)

    def decryptAttr(self, data, key):

        cbc = AES.new(self.getCipherKey(key), AES.MODE_CBC, "\0" * 16)
        attr = cbc.decrypt(self.b64_decode(data))
        self.logDebug("Decrypted Attr: " + attr)
        if not attr.startswith("MEGA"):
            self.fail(_("Decryption failed"))

        # Data is padded, 0-bytes must be stripped
        return json.loads(attr.replace("MEGA", "").rstrip("\0").strip())

    def decryptFile(self, key):
        """  Decrypts the file at lastDownload` """

        # upper 64 bit of counter start
        n = key[16:24]

        # convert counter to long and shift bytes
        ctr = Counter.new(128, initial_value=long(n.encode("hex"),16) << 64)
        cipher = AES.new(self.getCipherKey(key), AES.MODE_CTR, counter=ctr)

        self.pyfile.setStatus("decrypting")
        f = open(self.lastDownload, "rb")
        df = open(self.lastDownload.rstrip(self.FILE_SUFFIX), "wb")

        # TODO: calculate CBC-MAC for checksum

        size = 2 ** 15 # buffer size, 32k
        while True:
            buf = f.read(size)
            if not buf: break

            df.write(cipher.decrypt(buf))

        f.close()
        df.close()
        remove(self.lastDownload)

    def process(self, pyfile):

        key = None

        # match is guaranteed because plugin was chosen to handle url
        node = re.search(self.__pattern__, pyfile.url).group(2)
        if "!" in node:
            node, key = node.split("!")

        self.logDebug("File id: %s | Key: %s" % (node, key))

        if not key:
            self.fail(_("No file key provided in the URL"))

        # g is for requesting a download url
        # this is similar to the calls in the mega js app, documentation is very bad
        dl = self.callApi(a="g", g=1, p=node, ssl=1)[0]

        if "e" in dl:
            e = dl["e"]
            # ETEMPUNAVAIL (-18): Resource temporarily not available, please try again later
            if e == -18:
                self.retry()
            else:
                self.fail(_("Error code:") + e)

        # TODO: map other error codes, e.g
        # EACCESS (-11): Access violation (e.g., trying to write to a read-only share)

        key = self.b64_decode(key)
        attr = self.decryptAttr(dl["at"], key)

        pyfile.name = attr["n"] + self.FILE_SUFFIX

        self.download(dl["g"])
        self.decryptFile(key)

        # Everything is finished and final name can be set
        pyfile.name = attr["n"]

@@ -734,34 +734,15 @@
     
-return standard_b64decode(
+data = 
 data
@@ -781,13 +781,72 @@
 %22/%22)
-+ %22=%22
+%0A        return standard_b64decode(data + '=' * (-len(data) %25 4)
 )%0A%0A 

Support for incorrectly padded Mega.co.nz keys