Initial import.
[gitosis.git] / gitosis / ssh.py
1 import os, errno, re
2
3 def readKeys(keydir):
4     """
5     Read SSH public keys from ``keydir/*.pub``
6     """
7     for filename in os.listdir(keydir):
8         if filename.startswith('.'):
9             continue
10         basename, ext = os.path.splitext(filename)
11         if ext != '.pub':
12             continue
13
14         path = os.path.join(keydir, filename)
15         f = file(path)
16         try:
17             line = f.readline()
18         finally:
19             f.close()
20         line = line.rstrip('\n')
21         yield (basename, line)
22
23 COMMENT = '### autogenerated by gitosis, DO NOT EDIT'
24
25 def generateAuthorizedKeys(keys):
26     TEMPLATE=('command="gitosis-serve %(user)s",no-port-forwarding,'
27               +'no-X11-forwarding,no-agent-forwarding,no-pty %(key)s')
28
29     yield COMMENT
30     for (user, key) in keys:
31         yield TEMPLATE % dict(user=user, key=key)
32
33 _COMMAND_RE = re.compile('^command="(/[^ "]+/)?gitosis-serve [^"]+",no-port-forw'
34                          +'arding,no-X11-forwarding,no-agent-forwardi'
35                          +'ng,no-pty .*')
36
37 def filterAuthorizedKeys(fp):
38     """
39     Read lines from ``fp``, filter out autogenerated ones.
40
41     Note removes newlines.
42     """
43
44     for line in fp:
45         line = line.rstrip('\n')
46         if line == COMMENT:
47             continue
48         if _COMMAND_RE.match(line):
49             continue
50         yield line
51
52 def writeAuthorizedKeys(path, keydir):
53     tmp = '%s.%d.tmp' % (path, os.getpid())
54     try:
55         in_ = file(path)
56     except IOError, e:
57         if e.errno == errno.ENOENT:
58             in_ = None
59         else:
60             raise
61
62     try:
63         out = file(tmp, 'w')
64         try:
65             if in_ is not None:
66                 for line in filterAuthorizedKeys(in_):
67                     print >>out, line
68
69             keygen = readKeys(keydir)
70             for line in generateAuthorizedKeys(keygen):
71                 print >>out, line
72
73             os.fsync(out)
74         finally:
75             out.close()
76     finally:
77         if in_ is not None:
78             in_.close()
79     os.rename(tmp, path)
80
81 def _getParser():
82     import optparse
83     parser = optparse.OptionParser(
84         usage="%prog [--authkeys=FILE] KEYDIR")
85     parser.set_defaults(
86         authkeys=os.path.expanduser('~/.ssh/authorized_keys'),
87         )
88     parser.add_option(
89         "--authkeys",
90         help="path to SSH authorized keys file")
91     return parser
92
93 def main():
94     parser = _getParser()
95     (options, args) = parser.parse_args()
96
97     if len(args) != 1:
98         parser.error('Need one argument on the command line.')
99
100     keydir, = args
101
102     writeAuthorizedKeys(
103         path=options.authkeys,
104         keydir=keydir)