Add inform server

3 files changed, 83 insertions, 12 deletions

diff --git a/inform.py b/inform.py
index 1af39e6..6741415 100644
--- a/inform.py
+++ b/inform.py
@@ -1,4 +1,5 @@
 import json
+import copy
 import struct
 import binascii
 from Crypto.Cipher import AES
@@ -89,6 +90,17 @@ class InformPacket(object):
         self.data_version = None
         self.data_length = None
         self.raw_payload = None
+        self._used_key = None
+
+    def response_copy(self):
+        """Copy object for use in response
+
+        Generates a deep copy of the object and removes the payload so that it
+        can be used to respond to the station that send this inform request.
+        """
+        new_obj = copy.deepcopy(self)
+        new_obj.raw_payload = None
+        return new_obj
 
     @staticmethod
     def _format_mac_addr(mac_bytes):
@@ -112,10 +124,14 @@ class InformPacket(object):
     @property
     def payload(self):
         if self.data_version == 1:
-            return json.loads(self.raw_payload)
+            return json.loads(self.raw_payload.decode("latin-1"))
         else:
             return self.raw_payload
 
+    @payload.setter
+    def payload(self, value):
+        self.raw_payload = json.dumps(value)
+
 
 class InformSerializer(object):
     """Inform protocol version 1 parser/serializer
@@ -130,32 +146,41 @@ class InformSerializer(object):
     PROTOCOL_MAGIC = 1414414933
     MAX_VERSION = 1
 
-    def __init__(self, key=MASTER_KEY):
+    def __init__(self, key=None, key_bag=None):
         self.key = key
-        self._used_key = None
+        self.key_bag = key_bag or {}
 
     def _decrypt_payload(self, packet):
         if not packet.is_encrypted:
             return
 
-        for key in (self.key, self.MASTER_KEY):
+        i = 0
+        key = self.key_bag.get(packet.formatted_mac_addr)
+
+        for key in (key, self.key, self.MASTER_KEY):
+            if key is None:
+                continue
+
             decrypted = Cryptor(key, packet.iv).decrypt(packet.raw_payload)
 
             try:
-                json.loads(decrypted)
+                json.loads(decrypted.decode("latin-1"))
                 packet.raw_payload = decrypted
-                self._used_key = key
+                packet._used_key = key
                 break
-            except ValueError:
-                continue
+            except ValueError as err:
+                if err.message == "No JSON object could be decoded":
+                    continue
+                else:
+                    raise
 
     def parse(self, input):
         input_stream = BinaryDataStream(input)
 
         packet = InformPacket()
 
-        packet.magic = input_stream.read_int()
-        assert packet.magic == self.PROTOCOL_MAGIC
+        packet.magic_number = input_stream.read_int()
+        assert packet.magic_number == self.PROTOCOL_MAGIC
 
         packet.version = input_stream.read_int()
         assert packet.version < self.MAX_VERSION
@@ -175,13 +200,13 @@ class InformSerializer(object):
         if packet.data_version != 1:
             raise ValueError("Can no encrypt contents of pre 1.0 packets")
 
-        key = self._used_key if self._used_key else self.MASTER_KEY
+        key = packet._used_key if packet._used_key else self.MASTER_KEY
         return Cryptor(key, packet.iv).encrypt(json.dumps(packet.payload))
 
     def serialize(self, packet):
         output = BinaryDataStream.for_output()
 
-        output.write_int(packet.magic)
+        output.write_int(packet.magic_number)
         output.write_int(packet.version)
         output.write_string(packet.mac_addr)
         output.write_short(packet.flags)
diff --git a/inform_server.py b/inform_server.py
new file mode 100644
index 0000000..ebe3bda
--- /dev/null
+++ b/inform_server.py
@@ -0,0 +1,45 @@
+#!/usr/bin/python
+
+from StringIO import StringIO
+from SocketServer import TCPServer
+from inform import InformSerializer
+from ConfigParser import SafeConfigParser
+from SimpleHTTPServer import  SimpleHTTPRequestHandler
+
+
+class Handler(SimpleHTTPRequestHandler):
+
+    def __init__(self, *args, **kwargs):
+        super(SimpleHTTPServer, self).__init__(*args, **kwargs)
+        self.parser = InformSerializer(key_bag=get_keys("inform.cfg"))
+
+    def _get_keys(self, filename):
+        cfg = SafeConfigParser()
+        cfg.read(filename)
+
+        return dict((sect, cfg.get(sect, "authkey"))
+                for sect in cfg.sections())
+
+    def do_POST(self):
+        length = int(self.headers['content-length'])
+        body = StringIO(self.rfile.read(length))
+        packet = self.parser.parse(body)
+
+        noop_packet = packet.response_copy()
+        noop_packet.payload = { "_type": "noop", "interval": 10 }
+
+        buffer = StringIO(self.parser.serialize(noop_packet))
+
+        self.send_response(200)
+        self.send_header("Content-type", "application/x-binary")
+        self.send_header("Connection", "close")
+        self.send_header("User-Agent", "Unifi Controller")
+        self.send_header("Content-Length", buffer.len)
+        self.end_headers()
+
+        return buffer
+
+
+httpd = TCPServer(("", 9966), Handler)
+print "serving on 9966"
+httpd.serve_forever()
diff --git a/requirements.txt b/requirements.txt
index 953bf15..390e1fe 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1,2 @@
 tornado==3.1
+pycrypto==2.6