Convert existing SSH module to use the sshkey module, including all tests.

2 files changed, 20 insertions, 28 deletions

diff --git a/gitosis/ssh.py b/gitosis/ssh.py
index a9ed206..7b2c0c3 100644
--- a/gitosis/ssh.py
+++ b/gitosis/ssh.py
@@ -28,7 +28,7 @@ def readKeys(keydir):
         fp = file(path)
         for line in fp:
             line = line.rstrip('\n')
-            yield (basename, line)
+            yield (basename, sshkey.get_ssh_pubkey(line))
         fp.close()
 
 COMMENT = '### autogenerated by gitosis, DO NOT EDIT'
@@ -38,30 +38,14 @@ def generateAuthorizedKeys(keys):
     Genarate the lines for the Gitosis ~/.ssh/authorized_keys.
     """
     TEMPLATE = ('command="gitosis-serve %(user)s",no-port-forwarding,'
-                +'no-X11-forwarding,no-agent-forwarding,no-pty %(key)s')
+                +'no-X11-forwarding,no-agent-forwarding,no-pty %(key)s %(comment)s')
 
     yield COMMENT
     for (user, key) in keys:
-        yield TEMPLATE % dict(user=user, key=key)
+        yield TEMPLATE % dict(user=user, key=key.key, comment=key.comment)
 
-_COMMAND_OPTS_SAFE_CMD = \
-  'command="(/[^ "]+/)?gitosis-serve [^"]+"'
-_COMMAND_OPTS_SAFE = \
-  'no-port-forwarding' \
-+'|no-X11-forwarding' \
-+'|no-agent-forwarding' \
-+'|no-pty' \
-+'|from="[^"]*"'
-_COMMAND_OPTS_UNSAFE = \
-  'environment="[^"]*"' \
-+'|command="[^"]*"' \
-+'|permitopen="[^"]*"' \
-+'|tunnel="[^"]+"'
-
-_COMMAND_RE = re.compile(
-'^'+_COMMAND_OPTS_SAFE_CMD \
-+'(,('+_COMMAND_OPTS_SAFE+'))+' \
-+' .*')
+_GITOSIS_CMD_RE = '(/[^ "]+/)?gitosis-serve [^"]+'
+_COMMAND_RE = re.compile(_GITOSIS_CMD_RE)
 
 def filterAuthorizedKeys(fp):
     """
@@ -74,8 +58,13 @@ def filterAuthorizedKeys(fp):
         line = line.rstrip('\n')
         if line == COMMENT:
             continue
-        if _COMMAND_RE.match(line):
-            continue
+        try:
+            key = sshkey.get_ssh_pubkey(line)
+            if 'command' in key.options and \
+                _COMMAND_RE.match(key.options['command']):
+                continue
+        except sshkey.MalformedSSHKey:
+            pass
         yield line
 
 def writeAuthorizedKeys(path, keydir):
diff --git a/gitosis/test/test_ssh.py b/gitosis/test/test_ssh.py
index 77d7863..75effd5 100644
--- a/gitosis/test/test_ssh.py
+++ b/gitosis/test/test_ssh.py
@@ -4,6 +4,7 @@ import os
 from cStringIO import StringIO
 
 from gitosis import ssh
+from gitosis import sshkey
 from gitosis.test.util import mkdir, maketemp, writeFile, readFile
 
 def _key(s):
@@ -52,7 +53,9 @@ class ReadKeys_Test(object):
         writeFile(os.path.join(keydir, 'jdoe.pub'), KEY_1+'\n')
 
         gen = ssh.readKeys(keydir=keydir)
-        eq(gen.next(), ('jdoe', KEY_1))
+        (who, key) = gen.next()
+        eq(who, 'jdoe')
+        eq(key.full_key, KEY_1)
         assert_raises(StopIteration, gen.next)
 
     def test_two(self):
@@ -63,7 +66,7 @@ class ReadKeys_Test(object):
         writeFile(os.path.join(keydir, 'wsmith.pub'), KEY_2+'\n')
 
         gen = ssh.readKeys(keydir=keydir)
-        got = frozenset(gen)
+        got = frozenset( (i, j.full_key) for (i, j) in gen)
 
         eq(got,
            frozenset([
@@ -88,7 +91,7 @@ class ReadKeys_Test(object):
         writeFile(os.path.join(keydir, 'jdoe.pub'), KEY_1+'\n'+KEY_2+'\n')
 
         gen = ssh.readKeys(keydir=keydir)
-        got = frozenset(gen)
+        got = frozenset( (i, j.full_key) for (i, j) in gen)
 
         eq(got,
            frozenset([
@@ -99,8 +102,8 @@ class ReadKeys_Test(object):
 class GenerateAuthorizedKeys_Test(object):
     def test_simple(self):
         def k():
-            yield ('jdoe', KEY_1)
-            yield ('wsmith', KEY_2)
+            yield ('jdoe', sshkey.get_ssh_pubkey(KEY_1))
+            yield ('wsmith', sshkey.get_ssh_pubkey(KEY_2))
         gen = ssh.generateAuthorizedKeys(k())
         eq(gen.next(), ssh.COMMENT)
         eq(gen.next(), (