mirror-index/make-index.py

   1 #!/usr/bin/env python
   2 """make-index.py
   3
   4 Generates an nice index of the directories from a
   5 template.
   6
   7 Original Author: Jeremy Roman <jbroman@csclub.uwaterloo.ca>
   8
   9 So if you don't like how I did something,
  10 I'm the person you get to complain to.
  11 Please be gentle.
  12 """
  13
  14 import os, sys, time
  15 from subprocess import Popen, PIPE
  16 from optparse import OptionParser
  17 import yaml, mako.exceptions, webhelpers.html.tags
  18 from mako.template import Template
  19
  20 def reformat_size(size):
  21     """Reformats '124M' to '124 MB', et cetera."""
  22     if size[-1].isalpha():
  23         return size[:-1] + " " + size[-1] + "B"
  24     else:
  25         return size
  26
  27 def atomic_write(filename, body):
  28     """Atomically write to a file by writing a
  29     temporary file and then moving it to replace
  30     the desired output file.
  31
  32     This ensures that partial files are never seen
  33     by clients."""
  34
  35     # generate an appropriate temporary filename
  36     # in the same directory
  37     tmp_filename = "%s.%d.tmp" % (filename, os.getpid())
  38
  39     # open the directory so that we can fsync it
  40     dir = os.open(os.path.realpath(os.path.dirname(filename)), \
  41                       os.O_DIRECTORY | os.O_RDONLY)
  42
  43     # write to the temporary file
  44     tmp = open(tmp_filename, 'w')
  45     print >>tmp, body
  46     tmp.flush()
  47     os.fsync(tmp.fileno())
  48     tmp.close()
  49
  50     # atomically replace the actual file
  51     os.rename(tmp_filename, filename)
  52     os.fsync(dir)
  53     os.close(dir)
  54
  55 def main():
  56     # accept command-line arguments
  57     parser = OptionParser()
  58     parser.add_option("-c", "--config", dest="config", default="config.yaml",
  59                       help="configuration file to be used", metavar="FILE")
  60     parser.add_option("-D", "--docroot", dest="docroot",
  61                       help="directory to be scanned", metavar="DIR")
  62     parser.add_option("-F", "--duflags", dest="duflags",
  63                       help="flags to be passed to du, replaces any in config")
  64     parser.add_option("-o", "--output", dest="output", metavar="FILE",
  65                       help="file to which index page will be written. "
  66                       "Use /dev/stdout to send to standard out.")
  67     parser.add_option("-t", "--template", dest="template",
  68                       help="Mako template to render", metavar="FILE")
  69     parser.add_option("--nonatomic", dest="nonatomic", action="store_true",
  70                       default=False, help="write the output to the path "
  71                       "given without creating a temporary file in between. "
  72                       "This is automatically set if the output appears "
  73                       "to be a character device, not a file.")
  74     (options, args) = parser.parse_args()
  75
  76     # load config file
  77     try:
  78         config = yaml.load(file(options.config,'r'))
  79     except:
  80         config = None
  81
  82     if not config or type(config) != dict:
  83         print >>sys.stderr, "Unable to load configuration '%s'." % options.config
  84         sys.exit(-1)
  85
  86     # determine important variables based on an appropriate order of
  87     # precedence (command-line flags first, then the config file,
  88     # then built-in fallbacks)
  89     #
  90     # fallback value for nonatomic is used so that character devices
  91     # (e.g. /dev/stdout, /dev/null) are written to in the regular way
  92     docroot   = options.docroot   or config.get('docroot')
  93     duflags   = options.duflags   or config.get('duflags')   or "-h --max-depth=1"
  94     output    = options.output    or config.get('output')
  95     template  = options.template  or config.get("template")  or "index.mako"
  96     nonatomic = options.nonatomic or config.get("nonatomic") or \
  97         (os.path.exists(output) and not os.path.isfile(output))
  98
  99     # sanity checks
 100     if not docroot:
 101         print >>sys.stderr, "docroot not specified."
 102         print >>sys.stderr, "Define it in the config file or pass -D on the command line."
 103         sys.exit(-1)
 104     elif not output:
 105         print >>sys.stderr, "output not specified."
 106         print >>sys.stderr, "Define it in the config file or pass -o on the command line."
 107     elif not config.get('directories'):
 108         print >>sys.stderr, "directories not specified."
 109         print >>sys.stderr, "Define it in the config file."
 110         sys.exit(-1)
 111     elif not os.path.isdir(docroot):
 112         print >>sys.stderr, "docroot '%s' not found or not a directory." % docroot
 113         sys.exit(-1)
 114     elif not os.path.exists(template) or os.path.isdir(template):
 115         print >>sys.stderr, "template '%s' not found or is a directory." % template
 116         sys.exit(-1)
 117
 118     # Call du to compute size
 119     du = Popen(
 120         "/usr/bin/du %s %s | /usr/bin/sort -fk2" % (docroot, duflags),
 121         shell=True, stdout=PIPE, stderr=PIPE).communicate()
 122
 123     # Check that du executed successfully
 124     # If there's anything on stderr, send it
 125     # out our own stderr and terminate.
 126     if len(du[1].strip()) > 0:
 127         sys.stderr.write(du[1])
 128         print >>sys.stderr, "du terminated unsuccessfully. Not generating index."
 129         sys.exit(-1)
 130
 131     # first one should be total, grab its size and format
 132     du = du[0].splitlines() # we only care about stdout now
 133     total_size = reformat_size(du[0].split(None,2)[0])
 134
 135     # the rest are the sizes we want
 136     directories = []
 137     for line in du[1:]:
 138         (size, dir) = line.split(None, 2)
 139         dir = os.path.basename(dir)
 140         info = {'dir':dir, 'size':reformat_size(size)}
 141
 142         # use info from config.yaml, if found
 143         # otherwise, skip this directory
 144         if dir in config['directories']:
 145             info.update(config['directories'][dir])
 146         else:
 147             continue
 148
 149         directories.append(info)
 150
 151     # render the template to a string
 152     body = Template(filename=template).render(
 153         total_size=total_size,
 154         directories=directories,
 155         config=config,
 156         h=webhelpers.html.tags)
 157
 158     # write the rendered output
 159     if nonatomic:
 160         print >>file(output,'w'), body
 161     else:
 162         atomic_write(output, body)
 163
 164
 165 if __name__ == "__main__":
 166     main()