mirror of
https://github.com/mikf/gallery-dl.git
synced 2024-11-26 12:42:29 +01:00
205 lines
6.6 KiB
Python
205 lines
6.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
# Copyright 2015 Mike Fährmann
|
|
#
|
|
# This program is free software; you can redistribute it and/or modify
|
|
# it under the terms of the GNU General Public License version 2 as
|
|
# published by the Free Software Foundation.
|
|
|
|
import os
|
|
import sys
|
|
import re
|
|
import importlib
|
|
|
|
from .extractor.common import Message
|
|
|
|
class DownloadManager():
|
|
|
|
def __init__(self, opts, config):
|
|
self.opts = opts
|
|
self.config = config
|
|
self.modules = {}
|
|
self.extractors = ExtractorFinder(config)
|
|
|
|
def add(self, url):
|
|
job = DownloadJob(self, url)
|
|
job.run()
|
|
|
|
def get_downloader_module(self, scheme):
|
|
"""Return a downloader module suitable for 'scheme'"""
|
|
module = self.modules.get(scheme)
|
|
if module is None:
|
|
module = importlib.import_module(".downloader."+scheme, __package__)
|
|
self.modules[scheme] = module
|
|
return module
|
|
|
|
def get_base_directory(self):
|
|
"""Return the base-destination-directory for downloads"""
|
|
if self.opts.dest:
|
|
return self.opts.dest
|
|
else:
|
|
return self.config.get("general", "destination", fallback="/tmp/")
|
|
|
|
|
|
class DownloadJob():
|
|
|
|
def __init__(self, mngr, url):
|
|
self.mngr = mngr
|
|
self.extractor, self.info = mngr.extractors.get_for_url(url)
|
|
if self.extractor is None:
|
|
return
|
|
self.directory = mngr.get_base_directory()
|
|
self.downloaders = {}
|
|
self.filename_fmt = mngr.config.get(
|
|
self.info["category"], "filename",
|
|
fallback=self.info["filename"]
|
|
)
|
|
try:
|
|
segments = mngr.config.get(
|
|
self.info["category"], "directory"
|
|
).split("/")
|
|
except Exception:
|
|
segments = self.info["directory"]
|
|
self.directory_fmt = os.path.join(*segments)
|
|
|
|
def run(self):
|
|
"""Execute/Run the download job"""
|
|
if self.extractor is None:
|
|
return # TODO: error msg
|
|
|
|
for msg in self.extractor:
|
|
if msg[0] == Message.Url:
|
|
self.download(msg)
|
|
|
|
elif msg[0] == Message.Headers:
|
|
self.get_downloader("http:").set_headers(msg[1])
|
|
|
|
elif msg[0] == Message.Cookies:
|
|
self.get_downloader("http:").set_cookies(msg[1])
|
|
|
|
elif msg[0] == Message.Directory:
|
|
self.set_directory(msg)
|
|
|
|
elif msg[0] == Message.Version:
|
|
if msg[1] != 1:
|
|
raise "unsupported message-version ({}, {})".format(
|
|
self.info.category, msg[1]
|
|
)
|
|
# TODO: support for multiple message versions
|
|
|
|
def download(self, msg):
|
|
"""Download the resource specified in 'msg'"""
|
|
_, url, metadata = msg
|
|
filename = self.filename_fmt.format(**metadata)
|
|
path = os.path.join(self.directory, filename)
|
|
if os.path.exists(path):
|
|
self.print_skip(path)
|
|
return
|
|
downloader = self.get_downloader(url)
|
|
self.print_start(path)
|
|
tries = downloader.download(url, path)
|
|
self.print_success(path, tries)
|
|
|
|
def set_directory(self, msg):
|
|
"""Set and create the target directory for downloads"""
|
|
self.directory = os.path.join(
|
|
self.mngr.get_base_directory(),
|
|
self.directory_fmt.format(**msg[1])
|
|
)
|
|
os.makedirs(self.directory, exist_ok=True)
|
|
|
|
def get_downloader(self, url):
|
|
"""Return, and possibly construct, a downloader suitable for 'url'"""
|
|
pos = url.find(":")
|
|
scheme = url[:pos] if pos != -1 else "http"
|
|
if scheme == "https":
|
|
scheme = "http"
|
|
|
|
downloader = self.downloaders.get(scheme)
|
|
if downloader is None:
|
|
module = self.mngr.get_downloader_module(scheme)
|
|
downloader = module.Downloader()
|
|
self.downloaders[scheme] = downloader
|
|
|
|
return downloader
|
|
|
|
@staticmethod
|
|
def print_start(path):
|
|
"""Print a message indicating the start of a download"""
|
|
print(path, end="")
|
|
sys.stdout.flush()
|
|
|
|
@staticmethod
|
|
def print_skip(path):
|
|
"""Print a message indicating that a download has been skipped"""
|
|
print("\033[2m", path, "\033[0m", sep="")
|
|
|
|
@staticmethod
|
|
def print_success(path, tries):
|
|
"""Print a message indicating the completion of a download"""
|
|
if tries == 0:
|
|
print("\r", end="")
|
|
print("\r\033[1;32m", path, "\033[0m", sep="")
|
|
|
|
|
|
class ExtractorFinder():
|
|
|
|
def __init__(self, config):
|
|
self.config = config
|
|
|
|
def get_for_url(self, url):
|
|
"""Get an extractor-instance suitable for 'url'"""
|
|
name, match = self.find_pattern_match(url)
|
|
if match:
|
|
module = importlib.import_module(".extractor." + name, __package__)
|
|
klass = getattr(module, module.info["extractor"])
|
|
return klass(match, self.config), module.info
|
|
else:
|
|
print("no suitable extractor found")
|
|
return None, None
|
|
|
|
def find_pattern_match(self, url):
|
|
"""Find a pattern, that matches 'url', and return the (category,match) tuple"""
|
|
for category in self.config:
|
|
for key, value in self.config[category].items():
|
|
if key.startswith("regex"):
|
|
match = re.match(value, url)
|
|
if match:
|
|
return category, match
|
|
for category, info in self.extractor_metadata():
|
|
for pattern in info["pattern"]:
|
|
match = re.match(pattern, url)
|
|
if match:
|
|
return category, match
|
|
return None, None
|
|
|
|
def extractor_metadata(self):
|
|
"""Yield all extractor-name, -metadata tuples"""
|
|
path = os.path.join(os.path.dirname(__file__), "extractor")
|
|
for name in os.listdir(path):
|
|
extractor_path = os.path.join(path, name)
|
|
info = self.get_info_dict(extractor_path)
|
|
if info is not None:
|
|
yield os.path.splitext(name)[0], info
|
|
|
|
@staticmethod
|
|
def get_info_dict(extractor_path):
|
|
"""Get info-/metadata-dictionary for an extractor"""
|
|
try:
|
|
with open(extractor_path) as file:
|
|
for _ in range(30):
|
|
line = next(file)
|
|
if line.startswith("info ="):
|
|
break
|
|
else:
|
|
return None
|
|
|
|
info = [line[6:]]
|
|
for line in file:
|
|
info.append(line)
|
|
if line.startswith("}"):
|
|
break
|
|
except (StopIteration, OSError):
|
|
return None
|
|
return eval("".join(info))
|