Unifying public and private key interfaces

2 files changed, 29 insertions, 24 deletions

diff --git a/hgsshsign/__init__.py b/hgsshsign/__init__.py
index a6991db..b52dcc3 100644
--- a/hgsshsign/__init__.py
+++ b/hgsshsign/__init__.py
@@ -40,7 +40,7 @@ class SSHAuthority(object):
         private_key = ui.config("sshsign", "private_key", None)
         agent_socket = os.environ.get(SSHAgent.AGENT_SOCK_NAME)
         if private_key:
-            private_key = keys.load_private_key(absolute_path(private_key))
+            private_key = keys.PrivateKey.from_file(absolute_path(private_key))
         elif agent_socket:
             private_key = SSHAgent(agent_socket, key=public_key.blob)
         else:
@@ -55,16 +55,14 @@ class SSHAuthority(object):
 
     def verify(self, data, signature, whom):
         try:
-            key = self.key_manifest[whom] # XXX: More elegant error handling.
+            key = self.key_manifest[whom]
         except KeyError:
             raise util.Abort(_("No key found for %s" % whom))
 
         return key.verify(data, signature)
 
     def sign(self, data):
-        import hgsshsign.keys as keys
-
-        return keys.sign_like_agent(data, self.private_key)
+        return self.private_key.sign(data)
 
 
 def node2txt(repo, node, ver):
diff --git a/hgsshsign/keys.py b/hgsshsign/keys.py
index aa291eb..c723872 100644
--- a/hgsshsign/keys.py
+++ b/hgsshsign/keys.py
@@ -69,22 +69,29 @@ class PublicKey(object):
             fp.close()
 
 
-def load_private_key(filename):
-    fp = open(filename)
-    try:
-        first_line = fp.readline()
-    finally:
-        fp.close()
-
-    type = DSA if 'DSA' in first_line else RSA
-    return type.load_key(filename)
-
-
-def sign_like_agent(data, key):
-    """
-    Emulates the signing behavior of an ssh key agent.
-    """
-    digest = MessageDigest('sha1')
-    digest.update(data)
-    my_data = digest.final()
-    return key.sign(data)
+class PrivateKey(object):
+
+    def __init__(self, instance):
+        self.instance = instance
+
+    @classmethod
+    def from_file(cls, filename):
+        fp = open(filename)
+        try:
+            first_line = fp.readline()
+        finally:
+            fp.close()
+
+        type_ = DSA if 'DSA' in first_line else RSA
+        instance = type_.load_key(filename)
+
+        return cls(instance)
+
+    def sign(self, data):
+        """
+        Emulates the signing behavior of an ssh key agent.
+        """
+        digest = MessageDigest('sha1')
+        digest.update(data)
+        my_data = digest.final()
+        return self.instance.sign(data)