Merge pull request #45 from hbmartin/decrypt-performance

.gitignore

@@ -54,7 +54,11 @@ coverage.xml
 *.cover
 .hypothesis/
 .pytest_cache/
+*.mp4
+
+# Performance profiling
 prof/
+*.cprof
 
 # Debian Files
 debian/files

pytube/__main__.py

@@ -70,7 +70,6 @@ class YouTube:
         self.player_response: Dict = {}
         # streams
         self.age_restricted: Optional[bool] = None
-        self.vid_descr: Optional[str] = None
 
         self.fmt_streams: List[Stream] = []
 
@@ -125,8 +124,6 @@ class YouTube:
                 title = title[:index] if index > 0 else title
                 self.player_config_args["title"] = unescape(title)
 
-        if self.watch_html:
-            self.vid_descr = extract.get_vid_descr(self.watch_html)
         # https://github.com/nficano/pytube/issues/165
         stream_maps = ["url_encoded_fmt_stream_map"]
         if "adaptive_fmts" in self.player_config_args:
@@ -276,9 +273,9 @@ class YouTube:
         :rtype: str
 
         """
-        return self.vid_descr or (
-            self.player_response.get("videoDetails", {}).get("shortDescription")
-        )
+        return self.player_response.get("videoDetails", {}).get(
+            "shortDescription"
+        ) or extract.get_vid_descr(self.watch_html)
 
     @property
     def rating(self) -> float:

pytube/cipher.py

@@ -20,11 +20,75 @@ from itertools import chain
 from typing import List, Tuple, Dict, Callable, Any, Optional
 
 from pytube.exceptions import RegexMatchError
-from pytube.helpers import regex_search, create_logger
+from pytube.helpers import regex_search, create_logger, cache
 
 logger = create_logger()
 
 
+class Cipher:
+    def __init__(self, js: str):
+        self.transform_plan: List[str] = get_transform_plan(js)
+        var, _ = self.transform_plan[0].split(".")
+        self.transform_map = get_transform_map(js, var)
+        self.js_func_regex = re.compile(r"\w+\.(\w+)\(\w,(\d+)\)")
+
+    def get_signature(self, ciphered_signature: str) -> str:
+        """Decipher the signature.
+
+        Taking the ciphered signature, applies the transform functions.
+
+        :param str ciphered_signature:
+            The ciphered signature sent in the ``player_config``.
+        :rtype: str
+        :returns:
+           Decrypted signature required to download the media content.
+        """
+        signature = list(ciphered_signature)
+
+        for js_func in self.transform_plan:
+            name, argument = self.parse_function(js_func)  # type: ignore
+            signature = self.transform_map[name](signature, argument)
+            logger.debug(
+                "applied transform function\n"
+                "output: %s\n"
+                "js_function: %s\n"
+                "argument: %d\n"
+                "function: %s",
+                "".join(signature),
+                name,
+                argument,
+                self.transform_map[name],
+            )
+
+        return "".join(signature)
+
+    @cache
+    def parse_function(self, js_func: str) -> Tuple[str, int]:
+        """Parse the Javascript transform function.
+
+        Break a JavaScript transform function down into a two element ``tuple``
+        containing the function name and some integer-based argument.
+
+        :param str js_func:
+            The JavaScript version of the transform function.
+        :rtype: tuple
+        :returns:
+            two element tuple containing the function name and an argument.
+
+        **Example**:
+
+        >>> parse_function('DE.AJ(a,15)')
+        ('AJ', 15)
+
+        """
+        logger.debug("parsing transform function")
+        parse_match = self.js_func_regex.search(js_func)
+        if not parse_match:
+            raise RegexMatchError(caller="parse_function", pattern="js_func_regex")
+        fn_name, fn_arg = parse_match.groups()
+        return fn_name, int(fn_arg)
+
+
 def get_initial_function_name(js: str) -> str:
     """Extract the name of the function responsible for computing the signature.
     :param str js:
@@ -48,7 +112,6 @@ def get_initial_function_name(js: str) -> str:
         r"\bc\s*&&\s*[a-zA-Z0-9]+\.set\([^,]+\s*,\s*\([^)]*\)\s*\(\s*(?P<sig>[a-zA-Z0-9$]+)\(",  # noqa: E501
         r"\bc\s*&&\s*[a-zA-Z0-9]+\.set\([^,]+\s*,\s*\([^)]*\)\s*\(\s*(?P<sig>[a-zA-Z0-9$]+)\(",  # noqa: E501
     ]
-
     logger.debug("finding initial function name")
     for pattern in function_patterns:
         regex = re.compile(pattern)
@@ -71,7 +134,6 @@ def get_transform_plan(js: str) -> List[str]:
 
     **Example**:
 
-    >>> get_transform_plan(js)
     ['DE.AJ(a,15)',
     'DE.VR(a,3)',
     'DE.AJ(a,51)',
@@ -222,68 +284,3 @@ def map_functions(js_func: str) -> Callable:
         if re.search(pattern, js_func):
             return fn
     raise RegexMatchError(caller="map_functions", pattern="multiple")
-
-
-def parse_function(js_func: str) -> Tuple[str, int]:
-    """Parse the Javascript transform function.
-
-    Break a JavaScript transform function down into a two element ``tuple``
-    containing the function name and some integer-based argument.
-
-    :param str js_func:
-        The JavaScript version of the transform function.
-    :rtype: tuple
-    :returns:
-        two element tuple containing the function name and an argument.
-
-    **Example**:
-
-    >>> parse_function('DE.AJ(a,15)')
-    ('AJ', 15)
-
-    """
-    logger.debug("parsing transform function")
-    pattern = r"\w+\.(\w+)\(\w,(\d+)\)"
-    regex = re.compile(pattern)
-    parse_match = regex.search(js_func)
-    if not parse_match:
-        raise RegexMatchError(caller="parse_function", pattern=pattern)
-    fn_name, fn_arg = parse_match.groups()
-    return fn_name, int(fn_arg)
-
-
-def get_signature(js: str, ciphered_signature: str) -> str:
-    """Decipher the signature.
-
-    Taking the ciphered signature, applies the transform functions.
-
-    :param str js:
-        The contents of the base.js asset file.
-    :param str ciphered_signature:
-        The ciphered signature sent in the ``player_config``.
-    :rtype: str
-    :returns:
-       Decrypted signature required to download the media content.
-
-    """
-    transform_plan = get_transform_plan(js)
-    var, _ = transform_plan[0].split(".")
-    transform_map = get_transform_map(js, var)
-    signature = list(ciphered_signature)
-
-    for js_func in transform_plan:
-        name, argument = parse_function(js_func)
-        signature = transform_map[name](signature, argument)
-        logger.debug(
-            "applied transform function\n"
-            "output: %s\n"
-            "js_function: %s\n"
-            "argument: %d\n"
-            "function: %s",
-            "".join(signature),
-            name,
-            argument,
-            transform_map[name],
-        )
-
-    return "".join(signature)

pytube/extract.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 """This module contains all non-cipher related data extraction logic."""
 import json
-import pprint
 import re
 from collections import OrderedDict
 
@@ -10,7 +9,7 @@ from typing import Any, Optional, Tuple, List, Dict
 from urllib.parse import quote, parse_qs, unquote, parse_qsl
 from urllib.parse import urlencode
 
-from pytube import cipher
+from pytube.cipher import Cipher
 from pytube.exceptions import RegexMatchError, HTMLParseError, LiveStreamError
 from pytube.helpers import regex_search, logger
 
@@ -205,9 +204,10 @@ def get_ytplayer_config(html: str, age_restricted: bool = False) -> Any:
     return json.loads(yt_player_config)
 
 
-def get_vid_descr(html: str) -> str:
+def get_vid_descr(html: Optional[str]) -> str:
     html_parser = PytubeHTMLParser()
-    html_parser.feed(html)
+    if html:
+        html_parser.feed(html)
     return html_parser.vid_descr
 
 
@@ -224,6 +224,7 @@ def apply_signature(config_args: Dict, fmt: str, js: str) -> None:
         The contents of the base.js asset file.
 
     """
+    cipher = Cipher(js=js)
     stream_manifest = config_args[fmt]
     live_stream = (
         json.loads(config_args["player_response"])
@@ -247,17 +248,13 @@ def apply_signature(config_args: Dict, fmt: str, js: str) -> None:
             continue
 
         if js is not None:
-            signature = cipher.get_signature(js, stream["s"])
+            signature = cipher.get_signature(ciphered_signature=stream["s"])
         else:
             # signature not present in url (line 33), need js to descramble
             # TypeError caught in __main__
             raise TypeError("JS is None")
 
-        logger.debug(
-            "finished descrambling signature for itag=%s\n%s",
-            stream["itag"],
-            pprint.pformat({"s": stream["s"], "signature": signature,}, indent=2,),
-        )
+        logger.debug("finished descrambling signature for itag=%s", stream["itag"])
         # 403 forbidden fix
         stream_manifest[i]["url"] = url + "&sig=" + signature
 
@@ -320,6 +317,5 @@ def apply_descrambler(stream_data: Dict, key: str) -> None:
             {k: unquote(v) for k, v in parse_qsl(i)}
             for i in stream_data[key].split(",")
         ]
-    logger.debug(
-        "applying descrambler\n%s", pprint.pformat(stream_data[key], indent=2),
-    )
+
+    logger.debug("applying descrambler")

tests/test_cipher.py

@@ -20,17 +20,6 @@ def test_get_transform_object_with_no_match_should_error():
         cipher.get_transform_object("asdf", var="lt")
 
 
-def test_parse_function_with_match():
-    fn_name, fn_arg = cipher.parse_function("DE.AJ(a,15)")
-    assert fn_name == "AJ"
-    assert fn_arg == 15
-
-
-def test_parse_function_with_no_match_should_error():
-    with pytest.raises(RegexMatchError):
-        cipher.parse_function("asdf")
-
-
 def test_reverse():
     reversed_array = cipher.reverse([1, 2, 3, 4], None)
     assert reversed_array == [4, 3, 2, 1]

tests/test_streams.py

@@ -40,11 +40,11 @@ def test_title(cipher_signature):
 def test_description(cipher_signature):
     expected = (
         "PSY - ‘I LUV IT’ M/V @ https://youtu.be/Xvjnoagk6GU\n"
-        "PSY - ‘New Face’ M/V @https://youtu.be/OwJPPaEyqhI\n"
+        "PSY - ‘New Face’ M/V @https://youtu.be/OwJPPaEyqhI\n\n"
         "PSY - 8TH ALBUM '4X2=8' on iTunes @\n"
-        "https://smarturl.it/PSY_8thAlbum\n"
-        "PSY - GANGNAM STYLE(강남스타일) on iTunes @ http://smarturl.it/PsyGangnam\n"
-        "#PSY #싸이 #GANGNAMSTYLE #강남스타일\n"
+        "https://smarturl.it/PSY_8thAlbum\n\n"
+        "PSY - GANGNAM STYLE(강남스타일) on iTunes @ http://smarturl.it/PsyGangnam\n\n"
+        "#PSY #싸이 #GANGNAMSTYLE #강남스타일\n\n"
         "More about PSY@\nhttp://www.youtube.com/officialpsy\n"
         "http://www.facebook.com/officialpsy\n"
         "http://twitter.com/psy_oppa\n"
@@ -55,14 +55,14 @@ def test_description(cipher_signature):
     )
     assert cipher_signature.description == expected
 
-    cipher_signature.vid_descr = None
+    cipher_signature.player_response = {}
     expected = (
         "PSY - ‘I LUV IT’ M/V @ https://youtu.be/Xvjnoagk6GU\n"
-        "PSY - ‘New Face’ M/V @https://youtu.be/OwJPPaEyqhI\n\n"
+        "PSY - ‘New Face’ M/V @https://youtu.be/OwJPPaEyqhI\n"
         "PSY - 8TH ALBUM '4X2=8' on iTunes @\n"
-        "https://smarturl.it/PSY_8thAlbum\n\n"
-        "PSY - GANGNAM STYLE(강남스타일) on iTunes @ http://smarturl.it/PsyGangnam\n\n"
-        "#PSY #싸이 #GANGNAMSTYLE #강남스타일\n\n"
+        "https://smarturl.it/PSY_8thAlbum\n"
+        "PSY - GANGNAM STYLE(강남스타일) on iTunes @ http://smarturl.it/PsyGangnam\n"
+        "#PSY #싸이 #GANGNAMSTYLE #강남스타일\n"
         "More about PSY@\nhttp://www.youtube.com/officialpsy\n"
         "http://www.facebook.com/officialpsy\n"
         "http://twitter.com/psy_oppa\n"