code.delx.au - webdl/blob - autograbber.py

   1 #!/usr/bin/env python3
   2
   3 from common import load_root_node
   4 import fnmatch
   5 import logging
   6 import os
   7 import shutil
   8 import sys
   9
  10 HISTORY_FILENAME = ".history.txt"
  11 PATTERN_FILENAME = ".patterns.txt"
  12 EXCLUDE_FILENAME = ".excludes.txt"
  13
  14 class DownloadList(object):
  15     def __init__(self):
  16         self.exclude_list = set()
  17         self._load_exclude_list()
  18
  19         self.seen_list = set()
  20         self._load_history_file()
  21
  22         self.f = open(HISTORY_FILENAME, "a")
  23
  24     def _load_exclude_list(self):
  25         try:
  26             with open(EXCLUDE_FILENAME, "r") as f:
  27                 for line in f:
  28                     self.exclude_list.add(line.strip())
  29         except Exception as e:
  30             pass
  31
  32     def _load_history_file(self):
  33         self._move_old_file("downloaded_auto.txt")
  34         self._move_old_file(".downloaded_auto.txt")
  35
  36         try:
  37             with open(HISTORY_FILENAME, "r") as f:
  38                 for line in f:
  39                     self.seen_list.add(line.strip())
  40         except Exception as e:
  41             logging.error("Could not open history file: %s -- %s", HISTORY_FILENAME, e)
  42
  43     def _move_old_file(self, old_filename):
  44         if os.path.isfile(old_filename) and not os.path.isfile(HISTORY_FILENAME):
  45             logging.info("Migrating download history from %s to %s", old_filename, HISTORY_FILENAME)
  46             shutil.move(old_filename, HISTORY_FILENAME)
  47
  48     def wants(self, node):
  49         title = node.title.strip()
  50         if title in self.seen_list:
  51             return False
  52         for exclude in self.exclude_list:
  53             if fnmatch.fnmatch(title, exclude):
  54                 return False
  55         return True
  56
  57     def mark_seen(self, node):
  58         self.seen_list.add(node.title.strip())
  59         self.f.write(node.title.strip() + "\n")
  60         self.f.flush()
  61
  62
  63 def match(download_list, node, pattern, count=0):
  64     if node.can_download:
  65         if download_list.wants(node):
  66             if node.download():
  67                 download_list.mark_seen(node)
  68             else:
  69                 logging.error("Failed to download! %s", node.title)
  70         return
  71
  72     if count >= len(pattern):
  73         logging.error("No match found for pattern:", "/".join(pattern))
  74         return
  75     p = pattern[count]
  76     for child in node.get_children():
  77         if fnmatch.fnmatch(child.title, p):
  78             match(download_list, child, pattern, count+1)
  79
  80
  81 def process_one_dir(destdir, patternfile):
  82     os.chdir(destdir)
  83     node = load_root_node()
  84     download_list = DownloadList()
  85
  86     for line in open(patternfile):
  87         search = line.strip().split("/")
  88         match(download_list, node, search)
  89
  90 def check_directories(download_dirs):
  91     result = []
  92     failed = False
  93
  94     for d in download_dirs:
  95         d = os.path.abspath(d)
  96         if not os.path.isdir(d):
  97             print("Not a directory!", d)
  98             failed = True
  99
 100         pattern_filename = os.path.join(d, PATTERN_FILENAME)
 101         if not os.path.isfile(pattern_filename):
 102             print("Missing file!", pattern_filename)
 103             failed = True
 104
 105         result.append((d, pattern_filename))
 106
 107     if failed:
 108         print("Exiting!")
 109         sys.exit(1)
 110
 111     return result
 112
 113 def process_dirs(download_dirs):
 114     for download_dir, pattern_filename in check_directories(download_dirs):
 115         logging.info("Processing directory: %s", download_dir)
 116         process_one_dir(download_dir, pattern_filename)
 117
 118 if __name__ == "__main__":
 119     if len(sys.argv) <= 1:
 120         print("Usage: %s download_dir [download_dir ...]" % sys.argv[0])
 121         sys.exit(1)
 122
 123     if len(sys.argv) == 3 and os.path.isfile(sys.argv[2]):
 124         # Backwards compatibility with old argument format
 125         destdir = os.path.abspath(sys.argv[1])
 126         patternfile = os.path.abspath(sys.argv[2])
 127         run = lambda: process_one_dir(destdir, patternfile)
 128
 129     else:
 130         run = lambda: process_dirs(sys.argv[1:])
 131
 132     try:
 133         run()
 134     except (KeyboardInterrupt, EOFError):
 135         print("\nExiting...")