Fixing typos in agent and writing a quick integration test

4 files changed, 29 insertions, 21 deletions

diff --git a/.hgignore b/.hgignore
new file mode 100644
index 0000000..a26d142
--- /dev/null
+++ b/.hgignore
@@ -0,0 +1,2 @@
+syntax:glob
+*.pyc
diff --git a/ssh.py b/ssh.py
index a4209dc..4b7e94c 100644
--- a/ssh.py
+++ b/ssh.py
@@ -10,11 +10,12 @@ Commands to sign and verify revisions with your
 ssh key.
 """
 
+from structutils import bytes_to_int
+from sshagent import SSHAgent
 
+key = open('/Users/mcrute/.ssh/id_rsa.ag.pub').read()
+key = key.split()[1].decode('base64')
 
-
-
-
-
-if __name__ == '__main__':
-    pass
+agent = SSHAgent()
+signature = agent.sign("Hello world!", key)
+print bytes_to_int(signature)
diff --git a/sshagent.py b/sshagent.py
index 2a43f5f..32dc2f4 100644
--- a/sshagent.py
+++ b/sshagent.py
@@ -10,6 +10,7 @@ import os
 import socket
 import struct
 
+from structutils import int_to_bytes
 from structutils import pack_string, pack_int
 from structutils import unpack_int, unpack_string, unpack_mp_int
 
@@ -22,13 +23,14 @@ class SSHAgent(object):
     SSH2_AGENT_SIGN_RESPONSE = 14
     SSH2_AGENTC_SIGN_REQUEST = 13
 
-    def __init__(self, socket_path):
+    def __init__(self, socket_path=None):
         default_path = os.environ.get('SSH_AUTH_SOCK')
         socket_path = default_path if not socket_path else socket_path
 
         if not socket_path:
             raise ValueError("Could not find an ssh agent.")
 
+        self.socket_path = socket_path
         self.socket = None
 
     def connect(self):
@@ -43,16 +45,29 @@ class SSHAgent(object):
         to_send = ''.join([chr(SSHAgent.SSH2_AGENTC_SIGN_REQUEST),
                                key, data, flags])
         pkt_length = len(to_send)
-        packet = pack_int(pkg_length) + to_send
+        packet = pack_int(pkt_length) + to_send
 
         return packet
 
+    def sign(self, data, key):
+        if not self.socket:
+            self.connect()
+
+        packet = self._build_packet(data, key)
+
+        remaining = 0
+        while remaining < len(packet):
+            sent = self.socket.send(packet[remaining:])
+            remaining += sent
+
+        return self._parse_response()
+
     def _parse_response(self):
         response_length = unpack_int(self.socket.recv(4, socket.MSG_WAITALL))[0]
         if response_length == 1:
             raise ValueError("Agent failed")
 
-        response = auth_sock.recv(response_length, socket.MSG_WAITALL)
+        response = self.socket.recv(response_length, socket.MSG_WAITALL)
 
         status = ord(response[0])
         if status != SSHAgent.SSH2_AGENT_SIGN_RESPONSE:
@@ -62,14 +77,4 @@ class SSHAgent(object):
         _, remainder = unpack_string(remainder)
         response, _ = unpack_mp_int(remainder)
 
-        return response
-
-    def sign(self, data, key):
-        packet = self._build_packet(data, key)
-
-        remaining = 0
-        while remaining < len(packet):
-            sent = self.socket.send(packet[remaining:])
-            remaining += sent
-
-        return self._parse_response()
+        return int_to_bytes(response)
diff --git a/structutils.py b/structutils.py
index 26b0188..727d22e 100644
--- a/structutils.py
+++ b/structutils.py
@@ -75,7 +75,7 @@ def get_packed_mp_ints(buf, count=1):
     ints = []
     for _ in range(count):
         length, = INT_FORMAT.unpack(buf[:4])
-        ints.append(buf[:4+length]
+        ints.append(buf[:4+length])
         buf = buf[4+length:]
 
     return ints