Fixed missing import

[webdl] / common.py

diff --git a/common.py b/common.py
index da2051fcf9b7197e34161d6873ec510dcae10200..8b05719b2ff2cc177f0c1001311e065a6c6dc6e7 100644 (file)
--- a/common.py
+++ b/common.py
@@ -14,6 +14,7 @@ import subprocess
 import time
 import urllib.parse
 
+USER_AGENT = "Mozilla/5.0 (X11; Linux x86_64; rv:74.0) Gecko/20100101 Firefox/74.0"
 
 try:
     import autosocks
@@ -50,6 +51,7 @@ class Node(object):
     def get_children(self):
         if not self.children:
             self.fill_children()
+            self.children = natural_sort(self.children, key=lambda node: node.title)
         return self.children
 
     def fill_children(self):
@@ -88,7 +90,7 @@ def ensure_scheme(url):
     return urllib.parse.urlunparse(parts)
 
 http_session = requests.Session()
-http_session.headers["User-Agent"] = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.6; rv:21.0) Gecko/20100101 Firefox/21.0"
+http_session.headers["User-Agent"] = USER_AGENT
 
 def grab_text(url):
     logging.debug("grab_text(%r)", url)
@@ -149,25 +151,24 @@ def check_command_exists(cmd):
         return False
 
 def find_ffmpeg():
-    for ffmpeg in ["avconv", "ffmpeg"]:
-        if check_command_exists([ffmpeg, "--help"]):
-            return ffmpeg
+    if check_command_exists(["ffmpeg", "--help"]):
+        return "ffmpeg"
+
+    if check_command_exists(["avconv", "--help"]):
+        logging.warn("Detected libav-tools! ffmpeg is recommended")
+        return "avconv"
 
     raise Exception("You must install ffmpeg or libav-tools")
 
 def find_ffprobe():
-    for ffprobe in ["avprobe", "ffprobe"]:
-        if check_command_exists([ffprobe, "--help"]):
-            return ffprobe
-
-    raise Exception("You must install ffmpeg or libav-tools")
+    if check_command_exists(["ffprobe", "--help"]):
+        return "ffprobe"
 
-def find_streamlink():
-    for streamlink in ["streamlink", "livestreamer"]:
-        if check_command_exists([streamlink, "--help"]):
-            return streamlink
+    if check_command_exists(["avprobe", "--help"]):
+        logging.warn("Detected libav-tools! ffmpeg is recommended")
+        return "avprobe"
 
-    raise Exception("You must install streamlink or livestreamer")
+    raise Exception("You must install ffmpeg or libav-tools")
 
 def get_duration(filename):
     ffprobe = find_ffprobe()
@@ -180,12 +181,37 @@ def get_duration(filename):
     ]
     output = subprocess.check_output(cmd).decode("utf-8")
     for line in output.split("\n"):
-        if line.startswith("duration="):
-            return float(line.split("=")[1]) # ffprobe
-        if re.match(R'^[0-9.]*$', line):
-            return float(line) # avprobe
+        m = re.search(R"([0-9]+)", line)
+        if not m:
+            continue
+        duration = m.group(1)
+        if duration.isdigit():
+            return int(duration)
+
 
-    raise Exception("Unable to determine video duration of " + filename)
+    logging.debug("Falling back to full decode to find duration: %s % filename")
+
+    ffmpeg = find_ffmpeg()
+    cmd = [
+        ffmpeg,
+        "-i", filename,
+        "-vn",
+        "-f", "null", "-",
+    ]
+    output = subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8")
+    duration = None
+    for line in re.split(R"[\r\n]", output):
+        m = re.search(R"time=([0-9:]*)\.", line)
+        if not m:
+            continue
+        [h, m, s] = m.group(1).split(":")
+        # ffmpeg prints the duration as it reads the file, we want the last one
+        duration = int(h) * 3600 + int(m) * 60 + int(s)
+
+    if duration:
+        return duration
+    else:
+        raise Exception("Unable to determine video duration of " + filename)
 
 def check_video_durations(flv_filename, mp4_filename):
     flv_duration = get_duration(flv_filename)
@@ -240,8 +266,6 @@ def convert_to_mp4(filename):
 
 
 def download_hds(filename, video_url, pvswf=None):
-    streamlink = find_streamlink()
-
     filename = sanify_filename(filename)
     logging.info("Downloading: %s", filename)
 
@@ -252,9 +276,9 @@ def download_hds(filename, video_url, pvswf=None):
         param = video_url
 
     cmd = [
-        streamlink,
-        "-f",
-        "-o", filename,
+        "streamlink",
+        "--force",
+        "--output", filename,
         param,
         "best",
     ]
@@ -264,16 +288,15 @@ def download_hds(filename, video_url, pvswf=None):
         return False
 
 def download_hls(filename, video_url):
-    streamlink = find_streamlink()
-
     filename = sanify_filename(filename)
     video_url = "hlsvariant://" + video_url
     logging.info("Downloading: %s", filename)
 
     cmd = [
-        streamlink,
-        "-f",
-        "-o", filename,
+        "streamlink",
+        "--http-header", "User-Agent=" + USER_AGENT,
+        "--force",
+        "--output", filename,
         video_url,
         "best",
     ]
@@ -283,16 +306,14 @@ def download_hls(filename, video_url):
         return False
 
 def download_mpd(filename, video_url):
-    streamlink = find_streamlink()
-
     filename = sanify_filename(filename)
     video_url = "dash://" + video_url
     logging.info("Downloading: %s", filename)
 
     cmd = [
-        streamlink,
-        "-f",
-        "-o", filename,
+        "streamlink",
+        "--force",
+        "--output", filename,
         video_url,
         "best",
     ]