Example #1
0
    def test_requires_cycle(self):
        test1 = Argument("test1", requires="test2")
        test2 = Argument("test2", requires="test1")

        args = Arguments(test1, test2)

        self.assertRaises(RuntimeError, lambda: list(args.requires("test1")))
Example #2
0
    def test_requires_cycle(self):
        test1 = Argument("test1", requires="test2")
        test2 = Argument("test2", requires="test1")

        args = Arguments(test1, test2)

        self.assertRaises(RuntimeError, lambda: list(args.requires("test1")))
Example #3
0
    def test_requires(self):
        test1 = Argument("test1", requires="test2")
        test2 = Argument("test2", requires="test3")
        test3 = Argument("test3")

        args = Arguments(test1, test2, test3)

        self.assertEqual(list(args.requires("test1")), [test2, test3])
Example #4
0
    def test_getter(self):
        test1 = Argument("test1")
        test2 = Argument("test2")
        args = Arguments(test1, test2)

        self.assertEqual(args.get("test1"), test1)
        self.assertEqual(args.get("test2"), test2)
        self.assertEqual(args.get("test3"), None)
Example #5
0
    def test_requires(self):
        test1 = Argument("test1", requires="test2")
        test2 = Argument("test2", requires="test3")
        test3 = Argument("test3")

        args = Arguments(test1, test2, test3)

        self.assertEqual(list(args.requires("test1")), [test2, test3])
Example #6
0
    def test_getter(self):
        test1 = Argument("test1")
        test2 = Argument("test2")
        args = Arguments(test1, test2)

        self.assertEqual(args.get("test1"), test1)
        self.assertEqual(args.get("test2"), test2)
        self.assertEqual(args.get("test3"), None)
Example #7
0
    def test_setup_plugin_options(self):
        session = Mock()
        plugin = Mock(module="plugin")
        parser = argparse.ArgumentParser()
        parser.add_argument("--foo-foo", default=123)

        session.plugins = {"plugin": plugin}
        session.set_plugin_option = lambda name, key, value: session.plugins[name].options.update({key: value})
        plugin.arguments = Arguments(
            Argument("foo-foo", is_global=True),
            Argument("bar-bar", default=456),
            Argument("baz-baz", default=789, help=argparse.SUPPRESS)
        )

        with patch("streamlink_cli.main.args") as args:
            args.foo_foo = 321
            args.plugin_bar_bar = 654
            args.plugin_baz_baz = 987  # this wouldn't be set by the parser if the argument is suppressed

            setup_plugin_args(session, parser)
            self.assertEqual(plugin.options.get("foo_foo"), 123, "Sets the global-argument's default value")
            self.assertEqual(plugin.options.get("bar_bar"), 456, "Sets the plugin-argument's default value")
            self.assertEqual(plugin.options.get("baz_baz"), 789, "Sets the suppressed plugin-argument's default value")

            setup_plugin_options(session, plugin)
            self.assertEqual(plugin.options.get("foo_foo"), 321, "Sets the provided global-argument value")
            self.assertEqual(plugin.options.get("bar_bar"), 654, "Sets the provided plugin-argument value")
            self.assertEqual(plugin.options.get("baz_baz"), 789, "Doesn't set values of suppressed plugin-arguments")
Example #8
0
    def test_setup_plugin_args(self):
        session = Mock()
        plugin = Mock()
        parser = argparse.ArgumentParser()
        parser.add_argument("--global-arg1", default=123)
        parser.add_argument("--global-arg2", default=456)

        session.plugins = {"mock": plugin}
        plugin.arguments = Arguments(
            Argument("global-arg1", is_global=True),
            Argument("test1", default="default1"),
            Argument("test2", default="default2"),
            Argument("test3")
        )

        setup_plugin_args(session, parser)

        group = next((group for group in parser._action_groups if group.title == "Plugin options"), None)
        self.assertIsNotNone(group, "Adds the 'Plugin options' arguments group")
        self.assertEqual(
            [item for action in group._group_actions for item in action.option_strings],
            ["--mock-test1", "--mock-test2", "--mock-test3"],
            "Only adds plugin arguments and ignores global argument references"
        )

        self.assertEqual(plugin.options.get("global-arg1"), 123)
        self.assertEqual(plugin.options.get("global-arg2"), None)
        self.assertEqual(plugin.options.get("test1"), "default1")
        self.assertEqual(plugin.options.get("test2"), "default2")
        self.assertEqual(plugin.options.get("test3"), None)
Example #9
0
    def test_iter(self):
        test1 = Argument("test1")
        test2 = Argument("test2")
        args = Arguments(test1, test2)

        i_args = iter(args)

        self.assertEqual(next(i_args), test1)
        self.assertEqual(next(i_args), test2)
Example #10
0
    def test_set_defaults(self):
        session = Mock()
        plugin = Mock()
        parser = Mock()

        session.plugins = {"mock": plugin}
        plugin.arguments = Arguments(Argument("test1", default="default1"),
                                     Argument("test2", default="default2"),
                                     Argument("test3"))

        setup_plugin_args(session, parser)

        self.assertEqual(plugin.options.get("test1"), "default1")
        self.assertEqual(plugin.options.get("test2"), "default2")
        self.assertEqual(plugin.options.get("test3"), None)
Example #11
0
class Plugin(object):
    """A plugin can retrieve stream information from the URL specified.

    :param url: URL that the plugin will operate on
    """

    cache = None
    logger = None
    module = "unknown"
    options = Options()
    arguments = Arguments()
    session = None
    _user_input_requester = None

    @classmethod
    def bind(cls, session, module, user_input_requester=None):
        cls.cache = Cache(filename="plugin-cache.json", key_prefix=module)
        cls.logger = logging.getLogger("streamlink.plugin." + module)
        cls.module = module
        cls.session = session
        if user_input_requester is not None:
            if isinstance(user_input_requester, UserInputRequester):
                cls._user_input_requester = user_input_requester
            else:
                raise RuntimeError(
                    "user-input-requester must be an instance of UserInputRequester"
                )

    def __init__(self, url):
        self.url = url
        try:
            self.load_cookies()
        except RuntimeError:
            pass  # unbound cannot load

    @classmethod
    def can_handle_url(cls, url):
        raise NotImplementedError

    @classmethod
    def set_option(cls, key, value):
        cls.options.set(key, value)

    @classmethod
    def get_option(cls, key):
        return cls.options.get(key)

    @classmethod
    def get_argument(cls, key):
        return cls.arguments.get(key)

    @classmethod
    def stream_weight(cls, stream):
        return stream_weight(stream)

    @classmethod
    def default_stream_types(cls, streams):
        stream_types = ["rtmp", "hls", "hds", "http"]

        for name, stream in iterate_streams(streams):
            stream_type = type(stream).shortname()

            if stream_type not in stream_types:
                stream_types.append(stream_type)

        return stream_types

    @classmethod
    def broken(cls, issue=None):
        def func(*args, **kwargs):
            msg = (
                "This plugin has been marked as broken. This is likely due to "
                "changes to the service preventing a working implementation. ")

            if issue:
                msg += "More info: https://github.com/streamlink/streamlink/issues/{0}".format(
                    issue)

            raise PluginError(msg)

        def decorator(*args, **kwargs):
            return func

        return decorator

    @classmethod
    def priority(cls, url):
        """
        Return the plugin priority for a given URL, by default it returns
        NORMAL priority.
        :return: priority level
        """
        return NORMAL_PRIORITY

    def streams(self, stream_types=None, sorting_excludes=None):
        """Attempts to extract available streams.

        Returns a :class:`dict` containing the streams, where the key is
        the name of the stream, most commonly the quality and the value
        is a :class:`Stream` object.

        The result can contain the synonyms **best** and **worst** which
        points to the streams which are likely to be of highest and
        lowest quality respectively.

        If multiple streams with the same name are found, the order of
        streams specified in *stream_types* will determine which stream
        gets to keep the name while the rest will be renamed to
        "<name>_<stream type>".

        The synonyms can be fine tuned with the *sorting_excludes*
        parameter. This can be either of these types:

            - A list of filter expressions in the format
              *[operator]<value>*. For example the filter ">480p" will
              exclude streams ranked higher than "480p" from the list
              used in the synonyms ranking. Valid operators are >, >=, <
              and <=. If no operator is specified then equality will be
              tested.

            - A function that is passed to filter() with a list of
              stream names as input.


        :param stream_types: A list of stream types to return.
        :param sorting_excludes: Specify which streams to exclude from
                                 the best/worst synonyms.

        """

        try:
            ostreams = self._get_streams()
            if isinstance(ostreams, dict):
                ostreams = ostreams.items()

            # Flatten the iterator to a list so we can reuse it.
            if ostreams:
                ostreams = list(ostreams)
        except NoStreamsError:
            return {}
        except (IOError, OSError, ValueError) as err:
            raise PluginError(err)

        if not ostreams:
            return {}

        if stream_types is None:
            stream_types = self.default_stream_types(ostreams)

        # Add streams depending on stream type and priorities
        sorted_streams = sorted(iterate_streams(ostreams),
                                key=partial(stream_type_priority,
                                            stream_types))

        streams = {}
        for name, stream in sorted_streams:
            stream_type = type(stream).shortname()

            # Use * as wildcard to match other stream types
            if "*" not in stream_types and stream_type not in stream_types:
                continue

            # drop _alt from any stream names
            if name.endswith("_alt"):
                name = name[:-len("_alt")]

            existing = streams.get(name)
            if existing:
                existing_stream_type = type(existing).shortname()
                if existing_stream_type != stream_type:
                    name = "{0}_{1}".format(name, stream_type)

                if name in streams:
                    name = "{0}_alt".format(name)
                    num_alts = len(
                        list(
                            filter(lambda n: n.startswith(name),
                                   streams.keys())))

                    # We shouldn't need more than 2 alt streams
                    if num_alts >= 2:
                        continue
                    elif num_alts > 0:
                        name = "{0}{1}".format(name, num_alts + 1)

            # Validate stream name and discard the stream if it's bad.
            match = re.match("([A-z0-9_+]+)", name)
            if match:
                name = match.group(1)
            else:
                self.logger.debug(
                    "The stream '{0}' has been ignored "
                    "since it is badly named.", name)
                continue

            # Force lowercase name and replace space with underscore.
            streams[name.lower()] = stream

        # Create the best/worst synonmys
        def stream_weight_only(s):
            return (self.stream_weight(s)[0] or (len(streams) == 1 and 1))

        stream_names = filter(stream_weight_only, streams.keys())
        sorted_streams = sorted(stream_names, key=stream_weight_only)

        if isinstance(sorting_excludes, list):
            for expr in sorting_excludes:
                filter_func = stream_sorting_filter(expr, self.stream_weight)
                sorted_streams = list(filter(filter_func, sorted_streams))
        elif callable(sorting_excludes):
            sorted_streams = list(filter(sorting_excludes, sorted_streams))

        final_sorted_streams = OrderedDict()

        for stream_name in sorted(streams, key=stream_weight_only):
            final_sorted_streams[stream_name] = streams[stream_name]

        if len(sorted_streams) > 0:
            best = sorted_streams[-1]
            worst = sorted_streams[0]
            final_sorted_streams["worst"] = streams[worst]
            final_sorted_streams["best"] = streams[best]

        return final_sorted_streams

    def _get_streams(self):
        raise NotImplementedError

    def get_title(self):
        return None

    def get_author(self):
        return None

    def get_category(self):
        return None

    def save_cookies(self,
                     cookie_filter=None,
                     default_expires=60 * 60 * 24 * 7):
        """
        Store the cookies from ``http`` in the plugin cache until they expire. The cookies can be filtered
        by supplying a filter method. eg. ``lambda c: "auth" in c.name``. If no expiry date is given in the
        cookie then the ``default_expires`` value will be used.

        :param cookie_filter: a function to filter the cookies
        :type cookie_filter: function
        :param default_expires: time (in seconds) until cookies with no expiry will expire
        :type default_expires: int
        :return: list of the saved cookie names
        """
        if not self.session or not self.cache:
            raise RuntimeError("Cannot cache cookies in unbound plugin")

        cookie_filter = cookie_filter or (lambda c: True)
        saved = []

        for cookie in filter(cookie_filter, self.session.http.cookies):
            cookie_dict = {}
            for attr in ("version", "name", "value", "port", "domain", "path",
                         "secure", "expires", "discard", "comment",
                         "comment_url", "rfc2109"):
                cookie_dict[attr] = getattr(cookie, attr, None)
            cookie_dict["rest"] = getattr(cookie, "rest",
                                          getattr(cookie, "_rest", None))

            expires = default_expires
            if cookie_dict['expires']:
                expires = int(cookie_dict['expires'] - time.time())
            key = "__cookie:{0}:{1}:{2}:{3}".format(
                cookie.name, cookie.domain,
                cookie.port_specified and cookie.port or "80",
                cookie.path_specified and cookie.path or "*")
            self.cache.set(key, cookie_dict, expires)
            saved.append(cookie.name)

        if saved:
            self.logger.debug("Saved cookies: {0}".format(", ".join(saved)))
        return saved

    def load_cookies(self):
        """
        Load any stored cookies for the plugin that have not expired.

        :return: list of the restored cookie names
        """
        if not self.session or not self.cache:
            raise RuntimeError(
                "Cannot loaded cached cookies in unbound plugin")

        restored = []

        for key, value in self.cache.get_all().items():
            if key.startswith("__cookie"):
                cookie = requests.cookies.create_cookie(**value)
                self.session.http.cookies.set_cookie(cookie)
                restored.append(cookie.name)

        if restored:
            self.logger.debug("Restored cookies: {0}".format(
                ", ".join(restored)))
        return restored

    def clear_cookies(self, cookie_filter=None):
        """
        Removes all of the saved cookies for this Plugin. To filter the cookies that are deleted
        specify the ``cookie_filter`` argument (see :func:`save_cookies`).

        :param cookie_filter: a function to filter the cookies
        :type cookie_filter: function
        :return: list of the removed cookie names
        """
        if not self.session or not self.cache:
            raise RuntimeError(
                "Cannot loaded cached cookies in unbound plugin")

        cookie_filter = cookie_filter or (lambda c: True)
        removed = []

        for key, value in sorted(self.cache.get_all().items(),
                                 key=operator.itemgetter(0),
                                 reverse=True):
            if key.startswith("__cookie"):
                cookie = requests.cookies.create_cookie(**value)
                if cookie_filter(cookie):
                    del self.session.http.cookies[cookie.name]
                    self.cache.set(key, None, 0)
                    removed.append(key)

        return removed

    def input_ask(self, prompt):
        if self._user_input_requester:
            try:
                return self._user_input_requester.ask(prompt)
            except IOError as e:
                raise FatalPluginError("User input error: {0}".format(e))
            except NotImplementedError:  # ignore this and raise a FatalPluginError
                pass
        raise FatalPluginError(
            "This plugin requires user input, however it is not supported on this platform"
        )

    def input_ask_password(self, prompt):
        if self._user_input_requester:
            try:
                return self._user_input_requester.ask_password(prompt)
            except IOError as e:
                raise FatalPluginError("User input error: {0}".format(e))
            except NotImplementedError:  # ignore this and raise a FatalPluginError
                pass
        raise FatalPluginError(
            "This plugin requires user input, however it is not supported on this platform"
        )
Example #12
0
    def test_requires_invalid(self):
        test1 = Argument("test1", requires="test2")

        args = Arguments(test1)

        self.assertRaises(KeyError, lambda: list(args.requires("test1")))
Example #13
0
class Plugin:
    """
    Plugin base class for retrieving streams and metadata from the URL specified.
    """

    matchers: ClassVar[List[Matcher]] = None
    """
    The list of plugin matchers (URL pattern + priority).
    Use the :meth:`streamlink.plugin.pluginmatcher` decorator for initializing this list.
    """

    matches: Sequence[Optional[Match]]
    """A tuple of :class:`re.Match` results of all defined matchers"""

    matcher: Pattern
    """A reference to the compiled :class:`re.Pattern` of the first matching matcher"""

    match: Match
    """A reference to the :class:`re.Match` result of the first matching matcher"""

    # plugin metadata attributes
    id: Optional[str] = None
    """Metadata 'id' attribute: unique stream ID, etc."""
    title: Optional[str] = None
    """Metadata 'title' attribute: the stream's short descriptive title"""
    author: Optional[str] = None
    """Metadata 'author' attribute: the channel or broadcaster name, etc."""
    category: Optional[str] = None
    """Metadata 'category' attribute: name of a game being played, a music genre, etc."""

    cache = None
    logger = None
    module = "unknown"
    options = Options()
    arguments = Arguments()
    session = None
    _url: str = None
    _user_input_requester = None

    @classmethod
    def bind(cls, session, module, user_input_requester=None):
        cls.cache = Cache(filename="plugin-cache.json",
                          key_prefix=module)
        cls.logger = logging.getLogger("streamlink.plugins." + module)
        cls.module = module
        cls.session = session
        if user_input_requester is not None:
            if isinstance(user_input_requester, UserInputRequester):
                cls._user_input_requester = user_input_requester
            else:
                raise RuntimeError("user-input-requester must be an instance of UserInputRequester")

    @property
    def url(self) -> str:
        """
        The plugin's input URL.
        Setting a new value will automatically update the :attr:`matches`, :attr:`matcher` and :attr:`match` data.
        """

        return self._url

    @url.setter
    def url(self, value: str):
        self._url = value

        matches = [(pattern, pattern.match(value)) for pattern, priority in self.matchers or []]
        self.matches = tuple(m for p, m in matches)
        self.matcher, self.match = next(((p, m) for p, m in matches if m is not None), (None, None))

    def __init__(self, url: str):
        """
        :param str url: URL that the plugin will operate on
        """

        self.url = url

        try:
            self.load_cookies()
        except RuntimeError:
            pass  # unbound cannot load

    @classmethod
    def set_option(cls, key, value):
        cls.options.set(key, value)

    @classmethod
    def get_option(cls, key):
        return cls.options.get(key)

    @classmethod
    def get_argument(cls, key):
        return cls.arguments.get(key)

    @classmethod
    def stream_weight(cls, stream):
        return stream_weight(stream)

    @classmethod
    def default_stream_types(cls, streams):
        stream_types = ["hls", "http"]

        for name, stream in iterate_streams(streams):
            stream_type = type(stream).shortname()

            if stream_type not in stream_types:
                stream_types.append(stream_type)

        return stream_types

    @classmethod
    def broken(cls, issue=None):
        def func(*args, **kwargs):
            msg = (
                "This plugin has been marked as broken. This is likely due to "
                "changes to the service preventing a working implementation. "
            )

            if issue:
                msg += "More info: https://github.com/streamlink/streamlink/issues/{0}".format(issue)

            raise PluginError(msg)

        def decorator(*args, **kwargs):
            return func

        return decorator

    def streams(self, stream_types=None, sorting_excludes=None):
        """
        Attempts to extract available streams.

        Returns a :class:`dict` containing the streams, where the key is
        the name of the stream (most commonly the quality name), with the value
        being a :class:`Stream` instance.

        The result can contain the synonyms **best** and **worst** which
        point to the streams which are likely to be of highest and
        lowest quality respectively.

        If multiple streams with the same name are found, the order of
        streams specified in *stream_types* will determine which stream
        gets to keep the name while the rest will be renamed to
        "<name>_<stream type>".

        The synonyms can be fine-tuned with the *sorting_excludes*
        parameter, which can be one of these types:

            - A list of filter expressions in the format
              ``[operator]<value>``. For example the filter ">480p" will
              exclude streams ranked higher than "480p" from the list
              used in the synonyms ranking. Valid operators are ``>``, ``>=``, ``<``
              and ``<=``. If no operator is specified then equality will be tested.

            - A function that is passed to :meth:`filter` with a list of
              stream names as input.


        :param stream_types: A list of stream types to return
        :param sorting_excludes: Specify which streams to exclude from the best/worst synonyms
        :returns: A :class:`dict` of stream names and :class:`streamlink.stream.Stream` instances
        """

        try:
            ostreams = self._get_streams()
            if isinstance(ostreams, dict):
                ostreams = ostreams.items()

            # Flatten the iterator to a list so we can reuse it.
            if ostreams:
                ostreams = list(ostreams)
        except NoStreamsError:
            return {}
        except (OSError, ValueError) as err:
            raise PluginError(err)

        if not ostreams:
            return {}

        if stream_types is None:
            stream_types = self.default_stream_types(ostreams)

        # Add streams depending on stream type and priorities
        sorted_streams = sorted(iterate_streams(ostreams),
                                key=partial(stream_type_priority,
                                            stream_types))

        streams = {}
        for name, stream in sorted_streams:
            stream_type = type(stream).shortname()

            # Use * as wildcard to match other stream types
            if "*" not in stream_types and stream_type not in stream_types:
                continue

            # drop _alt from any stream names
            if name.endswith("_alt"):
                name = name[:-len("_alt")]

            existing = streams.get(name)
            if existing:
                existing_stream_type = type(existing).shortname()
                if existing_stream_type != stream_type:
                    name = "{0}_{1}".format(name, stream_type)

                if name in streams:
                    name = "{0}_alt".format(name)
                    num_alts = len(list(filter(lambda n: n.startswith(name), streams.keys())))

                    # We shouldn't need more than 2 alt streams
                    if num_alts >= 2:
                        continue
                    elif num_alts > 0:
                        name = "{0}{1}".format(name, num_alts + 1)

            # Validate stream name and discard the stream if it's bad.
            match = re.match("([A-z0-9_+]+)", name)
            if match:
                name = match.group(1)
            else:
                self.logger.debug(f"The stream '{name}' has been ignored since it is badly named.")
                continue

            # Force lowercase name and replace space with underscore.
            streams[name.lower()] = stream

        # Create the best/worst synonyms
        def stream_weight_only(s):
            return (self.stream_weight(s)[0] or (len(streams) == 1 and 1))

        stream_names = filter(stream_weight_only, streams.keys())
        sorted_streams = sorted(stream_names, key=stream_weight_only)
        unfiltered_sorted_streams = sorted_streams

        if isinstance(sorting_excludes, list):
            for expr in sorting_excludes:
                filter_func = stream_sorting_filter(expr, self.stream_weight)
                sorted_streams = list(filter(filter_func, sorted_streams))
        elif callable(sorting_excludes):
            sorted_streams = list(filter(sorting_excludes, sorted_streams))

        final_sorted_streams = OrderedDict()

        for stream_name in sorted(streams, key=stream_weight_only):
            final_sorted_streams[stream_name] = streams[stream_name]

        if len(sorted_streams) > 0:
            best = sorted_streams[-1]
            worst = sorted_streams[0]
            final_sorted_streams["worst"] = streams[worst]
            final_sorted_streams["best"] = streams[best]
        elif len(unfiltered_sorted_streams) > 0:
            best = unfiltered_sorted_streams[-1]
            worst = unfiltered_sorted_streams[0]
            final_sorted_streams["worst-unfiltered"] = streams[worst]
            final_sorted_streams["best-unfiltered"] = streams[best]

        return final_sorted_streams

    def _get_streams(self):
        """
        Implement the stream and metadata retrieval here.

        Needs to return either a dict of :class:`streamlink.stream.Stream` instances mapped by stream name, or needs to act
        as a generator which yields tuples of stream names and :class:`streamlink.stream.Stream` instances.
        """

        raise NotImplementedError

    def get_metadata(self) -> Dict[str, Optional[str]]:
        return dict(
            id=self.get_id(),
            author=self.get_author(),
            category=self.get_category(),
            title=self.get_title()
        )

    def get_id(self) -> Optional[str]:
        return None if self.id is None else str(self.id).strip()

    def get_title(self) -> Optional[str]:
        return None if self.title is None else str(self.title).strip()

    def get_author(self) -> Optional[str]:
        return None if self.author is None else str(self.author).strip()

    def get_category(self) -> Optional[str]:
        return None if self.category is None else str(self.category).strip()

    def save_cookies(
        self,
        cookie_filter: Optional[Callable] = None,
        default_expires: int = 60 * 60 * 24 * 7
    ) -> List[str]:
        """
        Store the cookies from :attr:`session.http` in the plugin cache until they expire. The cookies can be filtered
        by supplying a filter method. e.g. ``lambda c: "auth" in c.name``. If no expiry date is given in the
        cookie then the ``default_expires`` value will be used.

        :param cookie_filter: a function to filter the cookies
        :param default_expires: time (in seconds) until cookies with no expiry will expire
        :return: list of the saved cookie names
        """

        if not self.session or not self.cache:
            raise RuntimeError("Cannot cache cookies in unbound plugin")

        cookie_filter = cookie_filter or (lambda c: True)
        saved = []

        for cookie in filter(cookie_filter, self.session.http.cookies):
            cookie_dict = {}
            for attr in ("version", "name", "value", "port", "domain", "path", "secure", "expires", "discard",
                         "comment", "comment_url", "rfc2109"):
                cookie_dict[attr] = getattr(cookie, attr, None)
            cookie_dict["rest"] = getattr(cookie, "rest", getattr(cookie, "_rest", None))

            expires = default_expires
            if cookie_dict['expires']:
                expires = int(cookie_dict['expires'] - time.time())
            key = "__cookie:{0}:{1}:{2}:{3}".format(cookie.name,
                                                    cookie.domain,
                                                    cookie.port_specified and cookie.port or "80",
                                                    cookie.path_specified and cookie.path or "*")
            self.cache.set(key, cookie_dict, expires)
            saved.append(cookie.name)

        if saved:
            self.logger.debug("Saved cookies: {0}".format(", ".join(saved)))
        return saved

    def load_cookies(self) -> List[str]:
        """
        Load any stored cookies for the plugin that have not expired.

        :return: list of the restored cookie names
        """

        if not self.session or not self.cache:
            raise RuntimeError("Cannot load cached cookies in unbound plugin")

        restored = []

        for key, value in self.cache.get_all().items():
            if key.startswith("__cookie"):
                cookie = requests.cookies.create_cookie(**value)
                self.session.http.cookies.set_cookie(cookie)
                restored.append(cookie.name)

        if restored:
            self.logger.debug("Restored cookies: {0}".format(", ".join(restored)))
        return restored

    def clear_cookies(self, cookie_filter: Optional[Callable] = None) -> List[str]:
        """
        Removes all saved cookies for this plugin. To filter the cookies that are deleted
        specify the ``cookie_filter`` argument (see :func:`save_cookies`).

        :param cookie_filter: a function to filter the cookies
        :type cookie_filter: function
        :return: list of the removed cookie names
        """

        if not self.session or not self.cache:
            raise RuntimeError("Cannot clear cached cookies in unbound plugin")

        cookie_filter = cookie_filter or (lambda c: True)
        removed = []

        for key, value in sorted(self.cache.get_all().items(), key=operator.itemgetter(0), reverse=True):
            if key.startswith("__cookie"):
                cookie = requests.cookies.create_cookie(**value)
                if cookie_filter(cookie):
                    del self.session.http.cookies[cookie.name]
                    self.cache.set(key, None, 0)
                    removed.append(key)

        return removed

    def input_ask(self, prompt):
        if self._user_input_requester:
            try:
                return self._user_input_requester.ask(prompt)
            except OSError as e:
                raise FatalPluginError("User input error: {0}".format(e))
            except NotImplementedError:  # ignore this and raise a FatalPluginError
                pass
        raise FatalPluginError("This plugin requires user input, however it is not supported on this platform")

    def input_ask_password(self, prompt):
        if self._user_input_requester:
            try:
                return self._user_input_requester.ask_password(prompt)
            except OSError as e:
                raise FatalPluginError("User input error: {0}".format(e))
            except NotImplementedError:  # ignore this and raise a FatalPluginError
                pass
        raise FatalPluginError("This plugin requires user input, however it is not supported on this platform")
Example #14
0
    def test_requires_invalid(self):
        test1 = Argument("test1", requires="test2")

        args = Arguments(test1)

        self.assertRaises(KeyError, lambda: list(args.requires("test1")))