Daisuki plugin for livestreamer

daisuki.py

import base64
import json
import random
import re
import time

try:
    from Crypto import Random
    from Crypto.Cipher import AES, PKCS1_v1_5
    from Crypto.PublicKey import RSA
    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False

from livestreamer.compat import urljoin, urlparse, urlunparse
from livestreamer.exceptions import PluginError
from livestreamer.plugin import Plugin
from livestreamer.plugin.api import http, validate
from livestreamer.plugin.api.utils import parse_json
from livestreamer.stream import HLSStream

HDCORE_VERSION="3.2.0"

_url_re = re.compile(r"https?://www.daisuki.net/anime/watch/[^/]+/.+")
_flashvars_re = re.compile(r"var\s+flashvars\s*=\s*{([^}]*?)};", re.DOTALL)
_flashvar_re = re.compile(r"""(['"])(.*?)\1\s*:\s*(['"])(.*?)\3""")
_bgnwrapper_re = re.compile(r"""<script.*?src=(['"])(.*?bgnwrapper\.js.*?)\1""")

_schema = validate.Schema(
    validate.union({
        "flashvars": validate.all(
            validate.transform(_flashvars_re.search),
            validate.get(1),
            validate.transform(_flashvar_re.findall),
            validate.map(lambda v: (v[1], v[3])),
            validate.transform(dict),
            {
                "s": validate.text,
                "country": validate.text,
                "init": validate.text,
                validate.optional("ss_id"): validate.text,
                validate.optional("mv_id"): validate.text,
                validate.optional("device_cd"): validate.text,
                validate.optional("ss1_prm"): validate.text,
                validate.optional("ss2_prm"): validate.text,
                validate.optional("ss3_prm"): validate.text
            }
        ),
        "bgnwrapper": validate.all(
            validate.transform(_bgnwrapper_re.search),
            validate.get(2),
            validate.text
        )
    })
)

_language_schema = validate.Schema(
    validate.xml_findtext("./country_code")
)

_init_schema = validate.Schema(
    {
        "rtn": validate.all(
            validate.text
        )
    },
    validate.get("rtn")
)

def aes_encrypt(key, plaintext):
    aes = AES.new(key, AES.MODE_CBC, bytes([0] * AES.block_size))
    if len(plaintext) % AES.block_size != 0:
        plaintext += "\0" * (AES.block_size - len(plaintext) % AES.block_size)
    return base64.b64encode(aes.encrypt(plaintext))

def aes_decrypt(key, ciphertext):
    aes = AES.new(key, AES.MODE_CBC, bytes([0] * AES.block_size))
    plaintext = aes.decrypt(base64.b64decode(ciphertext))
    plaintext = plaintext.strip(b"\0")
    return plaintext.decode("utf-8")

def rsa_encrypt(key, plaintext):
    pubkey = RSA.importKey(key)
    cipher = PKCS1_v1_5.new(pubkey)
    return base64.b64encode(cipher.encrypt(plaintext))

def get_public_key(cache, url):
    headers = {}

    cached = cache.get("bgnwrapper")
    if cached and cached["url"] == url:
        headers["If-Modified-Since"] = cached["modified"]

    script = http.get(url, headers=headers)
    if cached and script.status_code == 304:
        return cached["pubkey"]

    modified = script.headers.get("Last-Modified", "")

    match = re.search(r"""(['"]-----BEGIN PUBLIC KEY-----.*?-----END PUBLIC KEY-----['"'])""", script.text, re.DOTALL)
    if match is None:
        return None

    matches = re.findall(r"""(['"])(.*?)\1""", match.group(1))
    if not matches:
        return

    pubkey = "\n".join([line[1].replace("\\n", "") for line in matches])

    cache.set("bgnwrapper", dict(url=url, modified=modified, pubkey=pubkey))

    return pubkey

class Daisuki(Plugin):
    @classmethod
    def can_handle_url(cls, url):
        return _url_re.match(url)

    def _get_streams(self):
        if not HAS_CRYPTO:
            raise PluginError("pyCrypto needs to be installed")

        page = http.get(self.url, schema=_schema)
        if not page:
            return

        pubkey_pem = get_public_key(self.cache, urljoin(self.url, page["bgnwrapper"]))
        if not pubkey_pem:
            raise PluginError("Unable to get public key")

        flashvars = page["flashvars"]

        params = {
            "cashPath":int(time.time()*1000)
        }
        res = http.get(urljoin(self.url, flashvars["country"]), params=params)
        if not res:
            return
        language = http.xml(res, schema=_language_schema)

        api_params = {}
        for key in ("ss_id", "mv_id", "device_cd", "ss1_prm", "ss2_prm", "ss3_prm"):
            if flashvars.get(key, ""):
                api_params[key] = flashvars[key]

        aeskey = bytes(random.getrandbits(8) for _ in range(32))

        data = {
            "s": flashvars["s"],
            "c": language,
            "e": self.url,
            "d": aes_encrypt(aeskey, json.dumps(api_params)),
            "a": rsa_encrypt(pubkey_pem, aeskey)
        }
        res = http.post(urljoin(self.url, flashvars["init"]), data=data)
        if not res:
            return
        rtn = http.json(res, schema=_init_schema)
        if not rtn:
            return

        init_data = parse_json(aes_decrypt(aeskey, rtn))
        
        parsed = urlparse(init_data["play_url"])
        if parsed.scheme != "http" or not parsed.path.startswith("/z/") or not parsed.path.endswith("/manifest.f4m"):
            return
        hlsstream_url = urlunparse(("https", parsed.netloc, "/i/" + parsed.path[3:-13] + "/master.m3u8", parsed.params, parsed.query, None))

        params = {
            "hdcore": HDCORE_VERSION,
        }
        return HLSStream.parse_variant_playlist(self.session, hlsstream_url, params=params)

__plugin__ = Daisuki