mercury.py

from __future__ import annotations
from librespot import util
from librespot.crypto import Packet
from librespot.proto import Mercury_pb2 as Mercury, Pubsub_pb2 as Pubsub
from librespot.structure import Closeable, PacketsReceiver, SubListener
from requests.structures import CaseInsensitiveDict
import io
import json
import logging
import queue
import struct
import threading
import typing

if typing.TYPE_CHECKING:
    from librespot.core import Session


class JsonMercuryRequest:
    request: RawMercuryRequest

    def __init__(self, request: RawMercuryRequest):
        self.request = request


class MercuryClient(Closeable, PacketsReceiver):
    logger = logging.getLogger("Librespot:MercuryClient")
    mercury_request_timeout = 3
    __callbacks: typing.Dict[int, Callback] = {}
    __remove_callback_lock = threading.Condition()
    __partials: typing.Dict[int, typing.List[bytes]] = {}
    __seq_holder = 0
    __seq_holder_lock = threading.Condition()
    __session: Session
    __subscriptions: typing.List[InternalSubListener] = []
    __subscriptions_lock = threading.Condition()

    def __init__(self, session: Session):
        self.__session = session

    def close(self) -> None:
        """
        Close the MercuryClient instance
        """
        if len(self.__subscriptions) != 0:
            for listener in self.__subscriptions:
                if listener.is_sub:
                    self.unsubscribe(listener.uri)
                else:
                    self.not_interested_in(listener.listener)
        if len(self.__callbacks) != 0:
            with self.__remove_callback_lock:
                self.__remove_callback_lock.wait(self.mercury_request_timeout)
        self.__callbacks.clear()

    def dispatch(self, packet: Packet) -> None:
        payload = io.BytesIO(packet.payload)
        seq_length = struct.unpack(">H", payload.read(2))[0]
        if seq_length == 2:
            seq = struct.unpack(">H", payload.read(2))[0]
        elif seq_length == 4:
            seq = struct.unpack(">i", payload.read(4))[0]
        elif seq_length == 8:
            seq = struct.unpack(">q", payload.read(8))[0]
        else:
            raise RuntimeError("Unknown seq length: {}".format(seq_length))
        flags = payload.read(1)
        parts = struct.unpack(">H", payload.read(2))[0]
        partial = self.__partials.get(seq)
        if partial is None or flags == 0:
            partial = []
            self.__partials[seq] = partial
        self.logger.debug(
            "Handling packet, cmd: 0x{}, seq: {}, flags: {}, parts: {}".format(
                util.bytes_to_hex(packet.cmd), seq, flags, parts))
        for _ in range(parts):
            size = struct.unpack(">H", payload.read(2))[0]
            buffer = payload.read(size)
            partial.append(buffer)
            self.__partials[seq] = partial
        if flags != b"\x01":
            return
        self.__partials.pop(seq)
        header = Mercury.Header()
        header.ParseFromString(partial[0])
        response = MercuryClient.Response(header, partial)
        if packet.is_cmd(Packet.Type.mercury_event):
            dispatched = False
            with self.__subscriptions_lock:
                for sub in self.__subscriptions:
                    if sub.matches(header.uri):
                        sub.dispatch(response)
                        dispatched = True
            if not dispatched:
                self.logger.debug(
                    "Couldn't dispatch Mercury event seq: {}, uri: {}, code: {}, payload: {}"
                    .format(seq, header.uri, header.status_code,
                            response.payload))
        elif (packet.is_cmd(Packet.Type.mercury_req)
              or packet.is_cmd(Packet.Type.mercury_sub)
              or packet.is_cmd(Packet.Type.mercury_sub)):
            callback = self.__callbacks.get(seq)
            self.__callbacks.pop(seq)
            if callback is not None:
                callback.response(response)
            else:
                self.logger.warning(
                    "Skipped Mercury response, seq: {}, uri: {}, code: {}".
                    format(seq, response.uri, response.status_code))
            with self.__remove_callback_lock:
                self.__remove_callback_lock.notify_all()
        else:
            self.logger.warning(
                "Couldn't handle packet, seq: {}, uri: {}, code: {}".format(
                    seq, header.uri, header.status_code))

    def interested_in(self, uri: str, listener: SubListener) -> None:
        self.__subscriptions.append(
            MercuryClient.InternalSubListener(uri, listener, False))

    def not_interested_in(self, listener: SubListener) -> None:
        try:
            for subscription in self.__subscriptions:
                if subscription.listener is listener:
                    self.__subscriptions.remove(subscription)
                    break
        except ValueError:
            pass

    def send(self, request: RawMercuryRequest, callback) -> int:
        """
        Send the Mercury request
        Args:
            request: RawMercuryRequest
            callback: Callback function
        Returns:
            MercuryClient.Response
        """
        buffer = io.BytesIO()
        seq: int
        with self.__seq_holder_lock:
            seq = self.__seq_holder
            self.__seq_holder += 1
        self.logger.debug(
            "Send Mercury request, seq: {}, uri: {}, method: {}".format(
                seq, request.header.uri, request.header.method))
        buffer.write(struct.pack(">H", 4))
        buffer.write(struct.pack(">i", seq))
        buffer.write(b"\x01")
        buffer.write(struct.pack(">H", 1 + len(request.payload)))
        header_bytes = request.header.SerializeToString()
        buffer.write(struct.pack(">H", len(header_bytes)))
        buffer.write(header_bytes)
        for part in request.payload:
            buffer.write(struct.pack(">H", len(part)))
            buffer.write(part)
        buffer.seek(0)
        cmd = Packet.Type.for_method(request.header.method)
        self.__session.send(cmd, buffer.read())
        self.__callbacks[seq] = callback
        return seq

    def send_sync(self, request: RawMercuryRequest) -> Response:
        """
        Send the Mercury request
        Args:
            request: RawMercuryRequest
        Returns:
            MercuryClient.Response
        """
        callback = MercuryClient.SyncCallback()
        seq = self.send(request, callback)
        try:
            response = callback.wait_response()
            if response is None:
                raise IOError(
                    "Request timeout out, {} passed, yet no response. seq: {}".
                    format(self.mercury_request_timeout, seq))
            return response
        except queue.Empty as e:
            raise IOError(e)

    def send_sync_json(self, request: JsonMercuryRequest) -> typing.Any:
        response = self.send_sync(request.request)
        if 200 <= response.status_code < 300:
            return json.loads(response.payload)
        raise MercuryClient.MercuryException(response)

    def subscribe(self, uri: str, listener: SubListener) -> None:
        """
        Subscribe URI
        Args:
            uri:
            listener:
        """
        response = self.send_sync(RawMercuryRequest.sub(uri))
        if response.status_code != 200:
            raise RuntimeError(response)
        if len(response.payload) > 0:
            for payload in response.payload:
                sub = Pubsub.Subscription()
                sub.ParseFromString(payload)
                self.__subscriptions.append(
                    MercuryClient.InternalSubListener(sub.uri, listener, True))
        else:
            self.__subscriptions.append(
                MercuryClient.InternalSubListener(uri, listener, True))
        self.logger.debug("Subscribed successfully to {}!".format(uri))

    def unsubscribe(self, uri) -> None:
        """
        Unsubscribe URI
        Args:
            uri:
        """
        response = self.send_sync(RawMercuryRequest.unsub(uri))
        if response.status_code != 200:
            raise RuntimeError(response)
        for subscription in self.__subscriptions:
            if subscription.matches(uri):
                self.__subscriptions.remove(subscription)
                break
        self.logger.debug("Unsubscribed successfully from {}!".format(uri))

    class Callback:
        def response(self, response: MercuryClient.Response) -> None:
            raise NotImplementedError

    class InternalSubListener:
        uri: str
        listener: SubListener
        is_sub: bool

        def __init__(self, uri: str, listener: SubListener, is_sub: bool):
            self.uri = uri
            self.listener = listener
            self.is_sub = is_sub

        def matches(self, uri: str) -> bool:
            """
            Compare with the URI given
            Args:
                uri: URI to be compared
            Returns:
                bool
            """
            return uri.startswith(self.uri)

        def dispatch(self, response: MercuryClient.Response) -> None:
            """
            Dispatch the event response
            Args:
                response: Response generated by the event
            """
            self.listener.event(response)

    class MercuryException(Exception):
        code: int

        def __init__(self, response: MercuryClient.Response):
            super().__init__("status: {}".format(response.status_code))
            self.code = response.status_code

    class PubSubException(MercuryException):
        pass

    class Response:
        uri: str
        payload: bytes
        status_code: int

        def __init__(self, header: Mercury.Header, payload: list[bytes]):
            self.uri = header.uri
            self.status_code = header.status_code
            self.payload = b"".join(payload[1:])

    class SyncCallback(Callback):
        __reference = queue.Queue()

        def response(self, response: MercuryClient.Response) -> None:
            """
            Set the response
            :param response:
            :return:
            """
            self.__reference.put(response)
            self.__reference.task_done()

        def wait_response(self) -> typing.Any:
            return self.__reference.get(
                timeout=MercuryClient.mercury_request_timeout)


class MercuryRequests:
    keymaster_client_id = "65b708073fc0480ea92a077233ca87bd"

    @staticmethod
    def get_root_playlists(username: str):
        """
        @TODO implement function
        """

    @staticmethod
    def request_token(device_id, scope):
        return JsonMercuryRequest(
            RawMercuryRequest.get(
                "hm://keymaster/token/authenticated?scope={}&client_id={}&device_id={}"
                .format(scope, MercuryRequests.keymaster_client_id,
                        device_id)))


class RawMercuryRequest:
    header: Mercury.Header
    payload: typing.List[bytes]

    def __init__(self, header: Mercury.Header, payload: typing.List[bytes]):
        self.header = header
        self.payload = payload

    @staticmethod
    def sub(uri: str):
        return RawMercuryRequest.new_builder().set_uri(uri).set_method(
            "SUB").build()

    @staticmethod
    def unsub(uri: str):
        return RawMercuryRequest.new_builder().set_uri(uri).set_method(
            "UNSUB").build()

    @staticmethod
    def get(uri: str):
        return RawMercuryRequest.new_builder().set_uri(uri).set_method(
            "GET").build()

    @staticmethod
    def send(uri: str, part: bytes):
        return RawMercuryRequest.new_builder().set_uri(uri) \
            .add_payload_part(part).set_method("SEND").build()

    @staticmethod
    def post(uri: str, part: bytes):
        return RawMercuryRequest.new_builder().set_uri(uri) \
            .set_method("POST").add_payload_part(part).build()

    @staticmethod
    def new_builder():
        return RawMercuryRequest.Builder()

    class Builder:
        header_dict: CaseInsensitiveDict
        payload: typing.List[bytes]

        def __init__(self):
            self.header_dict = CaseInsensitiveDict()
            self.payload = []

        def set_uri(self, uri: str):
            self.header_dict["uri"] = uri
            return self

        def set_content_type(self, content_type: str):
            self.header_dict["content_type"] = content_type
            return self

        def set_method(self, method: str):
            self.header_dict["method"] = method
            return self

        def add_user_field(self,
                           field: Mercury.UserField = None,
                           key: str = None,
                           value: str = None):
            if field is None and (key is None or value is None):
                return self
            try:
                self.header_dict["user_fields"]
            except KeyError:
                self.header_dict["user_fields"] = []
            if field is not None:
                self.header_dict["user_fields"].append(field)
            if key is not None and value is not None:
                self.header_dict["user_fields"].append(
                    Mercury.UserField(key=key, value=value.encode()))
            return self

        def add_payload_part(self, part: bytes):
            self.payload.append(part)
            return self

        def add_protobuf_payload(self, msg):
            return self.add_payload_part(msg)

        def build(self):
            return RawMercuryRequest(Mercury.Header(**self.header_dict),
                                     self.payload)