1 files changed, 92 insertions, 0 deletions
diff --git a/db_proxy.py b/db_proxy.py
new file mode 100644
index 0000000..f21b363
--- /dev/null
+++ b/db_proxy.py
@@ -0,0 +1,92 @@
+import sys
+import time
+import logging
+from threading import Event,Thread
+import Queue
+from proxy import OmegleProxyChat
+import pyPgSQL.PgSQL
+
+
+__all__ = []
+
+
+
+class DBThread(Thread):
+	def __init__(self,conn_str):
+		Thread.__init__(self)
+#        dbconn = ...("dbname='yvesf' user='dbuser' host='localhost' password='dbpass'")
+		self.conn = pyPgSQL.PgSQL.connect(conn_str)
+		self.queue = Queue.Queue()
+		self.running = True
+	def run(self):
+		#eat
+		while self.running:
+			try:
+				msg = self.queue.get(block=True,timeout=2)
+			except Queue.Empty:
+				continue
+			logging.getLogger("").debug("Insert %s" % msg)
+			cur = self.conn.cursor()
+			sql = """INSERT INTO omegle_messages(send_time,from_ident,to_ident,message)
+VALUES (CURRENT_TIMESTAMP, %s, %s, %s)"""
+			cur.execute(sql,(msg['from_ident'],msg['to_ident'],msg['message']))
+			cur.close()
+			self.conn.commit()
+
+	def add_message(self,from_ident,to_ident,message):
+		self.queue.put_nowait({'from_ident':from_ident,'to_ident':to_ident,'message':message})
+	def create_table(self):
+		cur=self.conn.cursor()
+		cur.execute("""
+CREATE TABLE omegle_messages(
+       send_time timestamp,
+       from_ident text,
+       to_ident text,
+       message text);""",())
+		cur.close()
+		self.conn.commit()
+	
+class DBOmegleProxyChat(OmegleProxyChat):
+	def __init__(self,disconnect_event,name,db,color=0):
+		OmegleProxyChat.__init__(self,disconnect_event,name,color)
+		self.db = db
+	def on_message(self,message):
+		OmegleProxyChat.on_message(self,message)
+		self.db.add_message(self.partner.id,self.id,message)
+
+if __name__ == "__main__":
+	colors = {"default":0, "black":30, "red":31, "green":32, "yellow":33,
+		  "blue":34,"magenta":35, "cyan":36, "white":37, "black":38, "black":39}
+
+	logging.basicConfig(level=logging.DEBUG)
+	
+	print "press ctrl-c to abort"
+	db=DBThread("")
+	db.start()
+	try:
+		db.create_table()
+	except Exception,e:
+		print e
+
+	event = Event()
+	A=DBOmegleProxyChat(event,"A",db,colors['red'])
+	B=DBOmegleProxyChat(event,"B",db,colors['blue'])
+
+	A.set_partner(B)
+	B.set_partner(A)
+
+	A.start()
+	B.start()
+
+	try:
+		while True:
+			event.wait(0.5)
+			if event.isSet():
+				print "One part disconnected, exit"
+				break
+	except KeyboardInterrupt:
+		print "CTRL-C pressed, exit"
+
+	if A.is_connected: A.disconnect()
+	if B.is_connected: B.disconnect()
+	db.running = False