Improved generation of filter subscription files

sitescripts/subscriptions/combineSubscriptions.py

 #!/usr/bin/env python
 # coding: utf-8
 
 # This file is part of the Adblock Plus web scripts,
 # Copyright (C) 2006-2013 Eyeo GmbH
 #
 # Adblock Plus is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License version 3 as
 # published by the Free Software Foundation.
 #
 # Adblock Plus is distributed in the hope that it will be useful,
 # but WITHOUT ANY WARRANTY; without even the implied warranty of
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU General Public License for more details.
 #
 # You should have received a copy of the GNU General Public License
 # along with Adblock Plus.  If not, see <http://www.gnu.org/licenses/>.
 
-import sys, os, re, subprocess, urllib2, time, traceback, codecs, hashlib, base64
+import sys, os, re, subprocess, urllib2, time, traceback, codecs, hashlib, base64, tempfile
 from getopt import getopt, GetoptError
 
 accepted_extensions = set([".txt"])
 ignore = set(["Apache.txt", "CC-BY-SA.txt", "GPL.txt", "MPL.txt"])
 verbatim = set(["COPYING"])
 
-def combine_subscriptions(sources, target_dir, timeout=30):
+def combine_subscriptions(sources, target_dir, timeout=30, tempdir=None):
   if not os.path.exists(target_dir):
     os.makedirs(target_dir, 0755)
+
+  def save_file(filename, data):
+    handle = tempfile.NamedTemporaryFile(mode="wb", dir=tempdir, delete=False)
+    handle.write(data.encode("utf-8"))
+    handle.close()
+
+    try:
+      subprocess.check_output(["7za", "a", "-tgzip", "-mx=9", "-bd", "-mpass=5", handle.name + ".gz", handle.name])
+    except:
+      print >>sys.stderr, "Failed to compress file %s. Please ensure that p7zip is installed on the system." % handle.name
+
+    path = os.path.join(target_dir, filename)
+    os.rename(handle.name, path)
+    os.rename(handle.name + ".gz", path + ".gz")
 
   known = set()
   for source_name, source in sources.iteritems():
     for filename in source.list_top_level_files():
       if filename in ignore or filename.startswith("."):
         continue
       if filename in verbatim:
-        process_verbatim_file(source, target_dir, filename)
+        process_verbatim_file(source, save_file, filename)
       elif not os.path.splitext(filename)[1] in accepted_extensions:
         continue
       else:
         try:
-          process_subscription_file(source_name, sources, target_dir, filename, timeout)
+          process_subscription_file(source_name, sources, save_file, filename, timeout)
         except:
           print >>sys.stderr, 'Error processing subscription file "%s"' % filename
           traceback.print_exc()
           print >>sys.stderr
         known.add(os.path.splitext(filename)[0] + ".tpl")
         known.add(os.path.splitext(filename)[0] + ".tpl.gz")
       known.add(filename)
       known.add(filename + ".gz")
 
   for filename in os.listdir(target_dir):
     if filename.startswith("."):
       continue
     if not filename in known:
       os.remove(os.path.join(target_dir, filename))
 
-def save_file(path, data):
-  handle = codecs.open(path, "wb", encoding="utf-8")
-  handle.write(data)
-  handle.close()
-  try:
-    subprocess.check_output(["7za", "a", "-tgzip", "-mx=9", "-bd", "-mpass=5", path + ".gz", path])
-  except:
-    print >>sys.stderr, "Failed to compress file %s. Please ensure that p7zip is installed on the system." % path
-
-def process_verbatim_file(source, target_dir, filename):
-  save_file(os.path.join(target_dir, filename), source.read_file(filename))
-
-def process_subscription_file(source_name, sources, target_dir, filename, timeout):
+def process_verbatim_file(source, save_file, filename):
+  save_file(filename, source.read_file(filename))
+
+def process_subscription_file(source_name, sources, save_file, filename, timeout):
   source = sources[source_name]
   lines = source.read_file(filename).splitlines()
 
   header = ""
   if len(lines) > 0:
     header = lines.pop(0)
   if not re.search(r"\[Adblock(?:\s*Plus\s*([\d\.]+)?)?\]", header, re.I):
     raise Exception("This is not a valid Adblock Plus subscription file.")
 
   lines = resolve_includes(source_name, sources, lines, timeout)
   seen = set(["checksum", "version"])
   def check_line(line):
     if line == "":
       return False
     match = re.search(r"^\s*!\s*(Redirect|Homepage|Title|Checksum|Version)\s*:", line, re.M | re.I)
     if not match:
       return True
     key = match.group(1).lower()
     if key in seen:
       return False
     seen.add(key)
     return True
   lines = filter(check_line, lines)
 
-  write_tpl(os.path.join(target_dir, os.path.splitext(filename)[0] + ".tpl"), lines)
+  write_tpl(save_file, os.path.splitext(filename)[0] + ".tpl", lines)
 
   lines.insert(0, "! Version: %s" % time.strftime("%Y%m%d%H%M", time.gmtime()))
 
   checksum = hashlib.md5()
   checksum.update("\n".join([header] + lines).encode("utf-8"))
   lines.insert(0, "! Checksum: %s" % base64.b64encode(checksum.digest()).rstrip("="))
   lines.insert(0, header)
-  save_file(os.path.join(target_dir, filename), "\n".join(lines))
+  save_file(filename, "\n".join(lines))
 
 def resolve_includes(source_name, sources, lines, timeout, level=0):
   if level > 5:
     raise Exception("There are too many nested includes, which is probably the result of a circular reference somewhere.")
 
   result = []
   for line in lines:
     match = re.search(r"^\s*%include\s+(.*)%\s*$", line)
     if match:
       filename = match.group(1)
@@ skipping 37 matching lines @@
       result.extend(newlines)
     else:
       if line.find("%timestamp%") >= 0:
         if level == 0:
           line = line.replace("%timestamp%", time.strftime("%d %b %Y %H:%M UTC", time.gmtime()))
         else:
           line = ""
       result.append(line)
   return result
 
-def write_tpl(path, lines):
+def write_tpl(save_file, filename, lines):
   result = []
   result.append("msFilterList")
   for line in lines:
     if re.search(r"^\s*!", line):
       # This is a comment. Handle "Expires" comment in a special way, keep the rest.
       match = re.search(r"\bExpires\s*(?::|after)\s*(\d+)\s*(h)?", line, re.I)
       if match:
         interval = int(match.group(1))
         if match.group(2):
           interval = int(interval / 24)
@@ skipping 78 matching lines @@
           line = line[2:]
         if domain:
           line = "%sd %s %s" % ("+" if is_exception else "-", domain, line)
           line = re.sub(r"\s+/$", "", line)
           result.append(line)
         elif is_exception:
           # Exception rules without domains are unsupported
           result.append("# " + origline)
         else:
           result.append("- " + line)
-  save_file(path, "\n".join(result) + "\n")
+  save_file(filename, "\n".join(result) + "\n")
 
 class FileSource:
   def __init__(self, dir):
     self._dir = dir
     if os.path.exists(os.path.join(dir, ".hg")):
       # This is a Mercurial repository, try updating
       subprocess.call(["hg", "-q", "-R", dir, "pull", "--update"])
 
   def get_path(self, filename):
     return os.path.join(self._dir, *filename.split("/"))
@@ skipping 40 matching lines @@
 
   timeout = 30
   for option, value in opts:
     if option in ("-h", "--help"):
       usage()
       sys.exit()
     elif option in ("-t", "--timeout"):
       timeout = int(value)
 
   combine_subscriptions(sources, target_dir, timeout)