allow rsa or dsa keys

2 files changed, 16 insertions, 5 deletions

diff --git a/server.py b/server.py
index 0137b2b..3bcfb3d 100644
--- a/server.py
+++ b/server.py
@@ -8,9 +8,9 @@ import ssh
 import termcast
 
 class Server(object):
-    def __init__(self, rsa_keyfile):
+    def __init__(self, keyfile):
         self.publisher = pubsub.Publisher()
-        self.rsa_keyfile = rsa_keyfile
+        self.keyfile = keyfile
 
     def listen(self):
         ssh_sock = self._open_socket(2200)
@@ -38,7 +38,7 @@ class Server(object):
     def handle_ssh_connection(self, client):
         self._handle_connection(
             client,
-            lambda client, connection_id: ssh.Connection(client, connection_id, self.publisher, self.rsa_keyfile)
+            lambda client, connection_id: ssh.Connection(client, connection_id, self.publisher, self.keyfile)
         )
 
     def handle_termcast_connection(self, client):
diff --git a/ssh.py b/ssh.py
index d4664bf..c320aa2 100644
--- a/ssh.py
+++ b/ssh.py
@@ -2,9 +2,20 @@ import paramiko
 import time
 
 class Connection(object):
-    def __init__(self, client, connection_id, publisher, rsa_keyfile):
+    def __init__(self, client, connection_id, publisher, keyfile):
         self.transport = paramiko.Transport(client)
-        self.transport.add_server_key(paramiko.RSAKey(filename=rsa_keyfile))
+
+        key = None
+        with open(keyfile) as f:
+            header = f.readline()
+            if header == "-----BEGIN DSA PRIVATE KEY-----\n":
+                key = paramiko.DSSKey(filename=keyfile)
+            elif header == "-----BEGIN RSA PRIVATE KEY-----\n":
+                key = paramiko.RSAKey(filename=keyfile)
+        if key is None:
+            raise Exception("%s doesn't appear to be an SSH keyfile" % keyfile)
+        self.transport.add_server_key(key)
+
         self.connection_id = connection_id
         self.publisher = publisher
         self.initialized = False