Initial import

2 files changed, 165 insertions, 0 deletions

diff --git a/redirecting_proxy.py b/redirecting_proxy.py
new file mode 100755
index 0000000..29f5df4
--- /dev/null
+++ b/redirecting_proxy.py
@@ -0,0 +1,84 @@
+#!/usr/bin/python
+import time
+import select
+import socket
+
+
+class Proxy:
+
+    buffer_size = 4096 * 2
+    delay = 0.0001
+
+    def __init__(self, host_port, to_host_port):
+        self.forward_to = to_host_port
+        self.input_list = []
+        self.channel = {}
+
+        self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        self.server.bind(host_port)
+        self.server.listen(200)
+
+    def main_loop(self):
+        self.input_list.append(self.server)
+
+        while time.sleep(self.delay) or True:
+            for self.s in select.select(self.input_list, [], [])[0]:
+                if self.s == self.server:
+                    self.on_accept()
+                    break
+
+                self.data = self.s.recv(self.buffer_size)
+                if len(self.data) == 0:
+                    self.on_close()
+                    break
+                else:
+                    self.on_recv()
+
+    def on_accept(self):
+        forward = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        clientsock, clientaddr = self.server.accept()
+        try:
+            forward.connect(self.forward_to)
+
+            self.input_list.append(clientsock)
+            self.input_list.append(forward)
+            self.channel[clientsock] = forward
+            self.channel[forward] = clientsock
+        except Exception:
+            clientsock.close()
+
+    def on_close(self):
+        out = self.channel[self.s]
+
+        self.input_list.remove(self.s)
+        self.input_list.remove(self.channel[self.s])
+
+        self.channel[out].close()
+        self.channel[self.s].close()
+
+        del self.channel[out]
+        del self.channel[self.s]
+
+    def on_recv(self):
+        if self.allow_data():
+            print repr(self.data)
+            self.channel[self.s].send(self.data)
+
+    def allow_data(self):
+        return not ("admin.$cmd" in self.data and "shutdown" in self.data)
+
+
+if __name__ == '__main__':
+    import os, sys
+
+    if len(sys.argv) < 3 or ":" not in sys.argv[2]:
+        program = os.path.basename(sys.argv[0])
+        print "usage: {} <port> <to_host>:<to_port>".format(program)
+        sys.exit(1)
+
+    try:
+        to_host, to_port = sys.argv[2].split(":")
+        Proxy(('', int(sys.argv[1])), (to_host, int(to_port))).main_loop()
+    except KeyboardInterrupt:
+        sys.exit(1)
diff --git a/translating_proxy.py b/translating_proxy.py
new file mode 100755
index 0000000..7e03e1b
--- /dev/null
+++ b/translating_proxy.py
@@ -0,0 +1,81 @@
+#!/usr/bin/env python
+
+import struct
+from bson import BSON
+from StringIO import StringIO
+
+
+DB_OFFSET = 16 + 4 # headers + flags
+
+OP_REPLY = 1
+OP_UPDATE = 2001
+OP_INSERT = 2002
+OP_QUERY = 2004
+OP_GET_MORE = 2005
+OP_DELETE = 2006
+
+# Touchup DB
+OP_TOUCHUP_DB = set((OP_UPDATE, OP_INSERT, OP_QUERY, OP_GET_MORE, OP_DELETE))
+
+def fixup_ns(raw_data, replace_with):
+    try:
+        data = BSON(raw_data).decode()
+    except:
+        return
+
+    if "ns" in data:
+        db, collection = data["ns"].split(".", 1)
+        data["ns"] = ".".join((replace_with, collection))
+        return BSON.encode(data)
+    else:
+        return raw_data
+
+
+def fixup_db(data, replace_with):
+    null = data.index("\x00")
+    db, collection = data[:null].split(".", 1)
+    replace_with = "admin" if db == "admin" else replace_with
+    fixed_up = "{}.{}\x00".format(replace_with, collection)
+    return fixed_up, (collection == "system.indexes")
+
+
+def handle_line(line):
+    output = StringIO()
+
+    length, req_id, resp_to, msg_type, flags = struct.unpack_from("<iiiii", line)
+
+    print msg_type
+
+    if msg_type not in OP_TOUCHUP_DB:
+        return
+
+    db_collection, rewrite_body = fixup_db(line[DB_OFFSET:], "mfi")
+
+    output.write(struct.pack("<iiii", req_id, resp_to, msg_type, flags))
+    output.write(db_collection)
+
+    if rewrite_body:
+        body_start = DB_OFFSET + len(db_collection) + 8
+        print repr(line[body_start:])
+        print BSON(line[body_start:]).decode()
+        output.write(fixup_ns(line[body_start:], "mfi"))
+    else:
+        body_start = DB_OFFSET + len(db_collection)
+        body = struct.pack("<i", 2280) + line[body_start+4:]
+        print struct.unpack_from("<i", line), len(body)
+        print repr(body)
+        print BSON(body).decode()
+        output.write(line[body_start:])
+
+    value = output.getvalue()
+#    print repr(struct.pack("<i", len(value)) + value)
+
+
+
+
+line = ''
+
+handle_line(line)
+
+#for line in open("unifi_snapshot.log"):
+#    handle_line(eval(line))