Commits

William Pitcock committed 1ff25c5

Add doctest tests.

  • Participants
  • Parent commits 78f1c70

Comments (0)

Files changed (1)

 This software is provided 'as is' and without any warranty, express or
 implied. In no event shall the authors be liable for any damages arising
 from the use of this software.
+
+The message module can be used to yield signed JSON structures, like so:
+
+>>> secret = '592783ded97d505759a8b25248c3fc5b'
+>>> message = { 'moo': 'cow' }
+>>> dumps(message, secret)
+'{"moo": "cow", "signature": "9117124d726d923719bf4a7faf15ef204c6da34fa3aa2279fc0b4f9ed1a6c2b2975fcf2bd520a8e92d3dc43264fcacc81937fdf58f579de42c9a5939a976b678"}'
 """
 
 import json
     """
     A class object which wraps a JSON structure providing signature
     and validation services.
+
+    An example of using the constructor:
+    >>> m = Message({ 'moo': 'cow' }, '592783ded97d505759a8b25248c3fc5b')
+    >>> m.payload_json()
+    '{"moo": "cow"}'
     """
     def __init__(self, obj, secret, expected_hash=None):
+        """
+        The constructor for the Message object.
+
+        Giving it an unsigned object should yield the following kind of object:
+
+        >>> m = Message({'moo': 'cow'}, '592783ded97d505759a8b25248c3fc5b')
+        >>> m.secret == '592783ded97d505759a8b25248c3fc5b'
+        True
+        >>> m.obj
+        {'moo': 'cow'}
+        >>> m.expected_hash
+        """
         self.obj = obj.copy()
         self.secret = secret
-        self.expected_hash = None
+        self.expected_hash = expected_hash
 
     def payload_json(self):
+        """
+        Returns the inner 'payload' JSON structure without the signature attached.
+
+        >>> m = Message({'moo': 'cow'}, '592783ded97d505759a8b25248c3fc5b')
+        >>> m.payload_json()
+        '{"moo": "cow"}'
+        """
         return json.dumps(self.obj)
 
     def signature(self):
+        """
+        Calculates the signature for the message.
+
+        The signature is just a PBKDF2 hash of the payload toasted against the secret, which is 64 bytes long.
+
+        >>> m = Message({'moo': 'cow'}, '592783ded97d505759a8b25248c3fc5b')
+        >>> m.signature()
+        '9117124d726d923719bf4a7faf15ef204c6da34fa3aa2279fc0b4f9ed1a6c2b2975fcf2bd520a8e92d3dc43264fcacc81937fdf58f579de42c9a5939a976b678'
+        """
         return pbkdf2_hex(self.payload_json(), self.secret, iterations=1000, keylen=64)
 
     def dumps(self, pretty_print=False):
+        """
+        Dumps the message with the signature calculated and attached.
+
+        >>> m = Message({'moo': 'cow'}, '592783ded97d505759a8b25248c3fc5b')
+        >>> m.dumps()
+        '{"moo": "cow", "signature": "9117124d726d923719bf4a7faf15ef204c6da34fa3aa2279fc0b4f9ed1a6c2b2975fcf2bd520a8e92d3dc43264fcacc81937fdf58f579de42c9a5939a976b678"}'
+        """
         envelope = self.obj
         envelope['signature'] = self.signature()
 
         return json.dumps(envelope)
 
     def validate(self):
+        """
+        Validates the message against the expected signature.  If the message is unsigned, this
+        always returns True.
+
+        >>> m = Message({'moo': 'cow'}, '592783ded97d505759a8b25248c3fc5b')
+        >>> m.validate()
+        True
+        >>> m = Message({'moo': 'cow'}, '592783ded97d505759a8b25248c3fc5b', '9117124d726d923719bf4a7faf15ef204c6da34fa3aa2279fc0b4f9ed1a6c2b2975fcf2bd520a8e92d3dc43264fcacc81937fdf58f579de42c9a5939a976b678')
+        >>> m.validate()
+        True
+        >>> m = Message({'moo': 'cow'}, '592783ded97d505759a8b25248c3fc5b', 'abcdef123456')
+        >>> m.validate()
+        False
+        """
         if self.expected_hash is None:
             return True
 
         return (self.signature() == self.expected_hash)
 
     def payload(self):
+        """
+        Returns the message payload.
+
+        >>> structure = {'moo': 'cow'}
+        >>> m = Message(structure, '592783ded97d505759a8b25248c3fc5b')
+        >>> m.payload()
+        {'moo': 'cow'}
+        >>> m.payload() == structure
+        True
+        """
         return self.obj
 
 class InvalidSignatureException(Exception):
     pass
 
-def loads(json_data, secret):
-    '''Unpack a JSON envelope containing an Edia RPC message.'''
+def loads(json_data, secret, allow_unsigned=False):
+    """
+    Unpack a JSON envelope containing an Edia RPC message.
+
+    >>> secret = '592783ded97d505759a8b25248c3fc5b'
+    >>> message = '{"moo": "cow", "signature": "9117124d726d923719bf4a7faf15ef204c6da34fa3aa2279fc0b4f9ed1a6c2b2975fcf2bd520a8e92d3dc43264fcacc81937fdf58f579de42c9a5939a976b678"}'
+    >>> loads(message, secret)
+    {u'moo': u'cow'}
+    >>> message = '{"moo": "cow", "signature": "abcdef123456"}'
+    >>> loads(message, secret)
+    Traceback (most recent call last):
+        ...
+    InvalidSignatureException
+    >>> message = '{"moo": "cow"}'
+    >>> loads(message, secret)
+    Traceback (most recent call last):
+        ...
+    InvalidSignatureException
+    >>> loads(message, secret, allow_unsigned=True)
+    {u'moo': u'cow'}
+    """
     envelope = json.loads(json_data)
 
     signature = envelope.pop('signature', None)
-    if signature is None:
+    if signature is None and allow_unsigned is not True:
         raise InvalidSignatureException()
 
     message = Message(envelope, secret, signature)
     return message.payload()
 
 def dumps(obj, secret, pretty_print=False):
-    '''Pack a dictionary into an Edia RPC message.'''
+    """
+    Pack a dictionary into an Edia RPC message.
+
+    >>> secret = '592783ded97d505759a8b25248c3fc5b'
+    >>> message = {'moo': 'cow'}
+    >>> signed_message = dumps(message, secret)
+    >>> signed_message
+    '{"moo": "cow", "signature": "9117124d726d923719bf4a7faf15ef204c6da34fa3aa2279fc0b4f9ed1a6c2b2975fcf2bd520a8e92d3dc43264fcacc81937fdf58f579de42c9a5939a976b678"}'
+    >>> signed_message == dumps(message, secret)
+    True
+    >>> unpack_message = loads(signed_message, secret)
+    >>> unpack_message == message
+    True
+    >>> signed_message == dumps(unpack_message, secret)
+    True
+    """
     return Message(obj, secret).dumps(pretty_print)
 
 if __name__ == '__main__':
-    secret = 'YNvngbwbWgjh4gbjh'
-    message = { 'message': 'doStuff', 'params': [1, 2, 3] }
-
-    print "Secret key is %s" % secret
-
-    jmsg = dumps(message, secret, True)
-    print "Signed JSON structure:", jmsg
-
-    print "Attempting to load signed JSON structure..."
-    unpack_message = loads(jmsg, secret)
-
-    print "Verifying that re-signing the structure yields the same structure..."
-    jmsg2 = dumps(unpack_message, secret, True)
-    assert jmsg == jmsg2
-
-    print "All tests complete."
+    import doctest
+    doctest.testmod()