summaryrefslogtreecommitdiff
path: root/db_proxy.py
blob: 3ae380ce69682adce7a4e7d1f605b1377bd4e1d9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
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,**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"