update config.yaml
[public/mirror.git] / 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()