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,**kwargs): Thread.__init__(self) self.conn = pyPgSQL.PgSQL.connect(client_encoding="utf-8", unicode_results=1,**kwargs) 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 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() logging.getLogger("").debug("commit" % msg) 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.INFO) print "press ctrl-c to abort" db=DBThread(**(dict(map(lambda x: x.split("="),sys.argv[1:])))) 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 print "main thread exit"