Completed
Push — master ( 0b3f02...a8c52a )
by Andrew
31s
created

WebRtcMessageHandler.close_file_receiver()   A

Complexity

Conditions 3

Size

Total Lines 8

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 3
c 1
b 0
f 0
dl 0
loc 8
rs 9.4285
1
import json
2
import logging
3
4
from django.core.exceptions import ValidationError
5
from django.db.models import Q
6
from tornado.gen import engine, Task
7
from tornadoredis import Client
8
9
from chat.log_filters import id_generator
10
from chat.models import Message, Room, RoomUsers
11
from chat.py2_3 import str_type
12
from chat.settings import ALL_ROOM_ID, SELECT_SELF_ROOM, TORNADO_REDIS_PORT, WEBRTC_CONNECTION
13
from chat.tornado.constants import VarNames, HandlerNames, Actions, RedisPrefix, WebRtcRedisStates
14
from chat.tornado.image_utils import process_images, prepare_img, save_images, get_message_images
15
from chat.tornado.message_creator import WebRtcMessageCreator
16
from chat.utils import get_max_key, execute_query, do_db, update_room, create_room_users, validate_edit_message, \
17
	get_or_create_room
18
19
parent_logger = logging.getLogger(__name__)
20
base_logger = logging.LoggerAdapter(parent_logger, {
21
	'id': 0,
22
	'ip': '000.000.000.000'
23
})
24
25
# TODO https://github.com/leporo/tornado-redis#connection-pool-support
26
# CONNECTION_POOL = tornadoredis.ConnectionPool(
27
# max_connections=500,
28
# wait_for_available=True)
29
30
31
class MessagesHandler(WebRtcMessageCreator):
32
33
	def __init__(self, *args, **kwargs):
34
		self.closed_channels = None
35
		self.parsable_prefix = 'p'
36
		super(MessagesHandler, self).__init__(*args, **kwargs)
37
		self.webrtc_ids = {}
38
		self.ip = None
39
		from chat import global_redis
40
		self.async_redis_publisher = global_redis.async_redis_publisher
41
		self.sync_redis = global_redis.sync_redis
42
		self.channels = []
43
		self._logger = None
44
		self.async_redis = Client(port=TORNADO_REDIS_PORT)
45
		self.patch_tornadoredis()
46
		self.pre_process_message = {
47
			Actions.GET_MESSAGES: self.process_get_messages,
48
			Actions.SEND_MESSAGE: self.process_send_message,
49
			Actions.CREATE_DIRECT_CHANNEL: self.create_user_channel,
50
			Actions.DELETE_ROOM: self.delete_channel,
51
			Actions.EDIT_MESSAGE: self.edit_message,
52
			Actions.CREATE_ROOM_CHANNEL: self.create_new_room,
53
			Actions.INVITE_USER: self.invite_user
54
		}
55
		self.post_process_message = {
56
			Actions.CREATE_DIRECT_CHANNEL: self.send_client_new_channel,
57
			Actions.CREATE_ROOM_CHANNEL: self.send_client_new_channel,
58
			Actions.DELETE_ROOM: self.send_client_delete_channel,
59
			Actions.INVITE_USER: self.send_client_new_channel
60
		}
61
62
	def patch_tornadoredis(self):  # TODO remove this
63
		fabric = type(self.async_redis.connection.readline)
64
		self.async_redis.connection.old_read = self.async_redis.connection.readline
65
66
		def new_read(new_self, callback=None):
67
			try:
68
				return new_self.old_read(callback=callback)
69
			except Exception as e:
70
				current_online = self.get_online_from_redis(RedisPrefix.DEFAULT_CHANNEL)
71
				self.logger.error(e)
72
				self.logger.error(
73
					"Exception info: "
74
					"self.id: %s ;;; "
75
					"self.connected = '%s';;; "
76
					"Redis default channel online = '%s';;; "
77
					"self.channels = '%s';;; "
78
					"self.closed_channels  = '%s';;;",
79
					self.id, self.connected, current_online, self.channels, self.closed_channels
80
				)
81
				raise e
82
83
		self.async_redis.connection.readline = fabric(new_read, self.async_redis.connection)
84
85
	@property
86
	def connected(self):
87
		raise NotImplemented
88
89
	@connected.setter
90
	def connected(self, value):
91
		raise NotImplemented
92
93
	@engine
94
	def listen(self, channels):
95
		yield Task(
96
			self.async_redis.subscribe, channels)
97
		self.async_redis.listen(self.pub_sub_message)
98
99
	@property
100
	def logger(self):
101
		return self._logger if self._logger else base_logger
102
103
	@engine
104
	def add_channel(self, channel):
105
		self.channels.append(channel)
106
		yield Task(
107
			self.async_redis.subscribe, (channel,))
108
109
	@staticmethod
110
	def evaluate(query_set):
111
		do_db(len, query_set)
112
		return query_set
113
114
	def get_online_from_redis(self, channel):
115
		return self.get_online_and_status_from_redis(channel)[1]
116
117
	def get_online_and_status_from_redis(self, channel):
118
		"""
119
		:rtype : (bool, list)
120
		"""
121
		online = self.sync_redis.ssmembers(channel)
122
		self.logger.debug('!! channel %s redis online: %s', channel, online)
123
		return self.parse_redis_online(online) if online else (False, [])
124
125
	def parse_redis_online(self, online):
126
		"""
127
		:rtype : (bool, list)
128
		"""
129
		result = set()
130
		user_is_online = False
131
		for decoded in online:  # py2 iteritems
132
			# : char specified in cookies_middleware.py.create_id
133
			user_id = int(decoded.split(':')[0])
134
			if user_id == self.user_id and decoded != self.id:
135
				user_is_online = True
136
			result.add(user_id)
137
		return user_is_online, list(result)
138
139
	def add_online_user(self, room_id, offline_messages=None):
140
		"""
141
		adds to redis
142
		online_users = { connection_hash1 = stored_redis_user1, connection_hash_2 = stored_redis_user2 }
143
		:return:
144
		"""
145
		self.async_redis_publisher.sadd(room_id, self.id)
146
		# since we add user to online first, latest trigger will always show correct online
147
		is_online, online = self.get_online_and_status_from_redis(room_id)
148
		if is_online:  # Send user names to self
149
			online_user_names_mes = self.room_online(online, Actions.REFRESH_USER, room_id)
150
			self.logger.info('!! Second tab, retrieving online for self')
151
			self.ws_write(online_user_names_mes)
152
		else:  # if a new tab has been opened
153
			online.append(self.user_id)
154
			online_user_names_mes = self.room_online(online, Actions.LOGIN, room_id)
155
			self.logger.info('!! First tab, sending refresh online for all')
156
			self.publish(online_user_names_mes, room_id)
157
			if offline_messages:
158
				self.ws_write(self.load_offline_message(offline_messages, room_id))
159
160
	def publish(self, message, channel, parsable=False):
161
		jsoned_mess = json.dumps(message)
162
		self.logger.debug('<%s> %s', channel, jsoned_mess)
163
		if parsable:
164
			jsoned_mess = self.encode(jsoned_mess)
165
		self.async_redis_publisher.publish(channel, jsoned_mess)
166
167
	def encode(self, message):
168
		"""
169
		Marks message with prefix to specify that
170
		it should be decoded and proccesed before sending to client
171
		@param message: message to mark
172
		@return: marked message
173
		"""
174
		return self.parsable_prefix + message
175
176
	def remove_parsable_prefix(self, message):
177
		if message.startswith(self.parsable_prefix):
178
			return message[1:]
179
180
	def pub_sub_message(self, message):
181
		data = message.body
182
		if isinstance(data, str_type):  # subscribe event
183
			prefixless_str = self.remove_parsable_prefix(data)
184
			if prefixless_str:
185
				dict_message = json.loads(prefixless_str)
186
				res = self.post_process_message[dict_message[VarNames.EVENT]](dict_message)
187
				if not res:
188
					self.ws_write(prefixless_str)
189
			else:
190
				self.ws_write(data)
191
192
	def ws_write(self, message):
193
		raise NotImplementedError('WebSocketHandler implements')
194
195
	def process_send_message(self, message):
196
		"""
197
		:type message: dict
198
		"""
199
		raw_imgs = message.get(VarNames.IMG)
200
		channel = message[VarNames.CHANNEL]
201
		message_db = Message(
202
			sender_id=self.user_id,
203
			content=message[VarNames.CONTENT],
204
			symbol=get_max_key(raw_imgs)
205
		)
206
		message_db.room_id = channel
207
		do_db(message_db.save)
208
		db_images = save_images(raw_imgs, message_db.id)
209
		prepared_message = self.create_send_message(
210
			message_db,
211
			Actions.PRINT_MESSAGE,
212
			prepare_img(db_images, message_db.id)
213
		)
214
		self.publish(prepared_message, channel)
215
216
	def create_new_room(self, message):
217
		room_name = message[VarNames.ROOM_NAME]
218
		if not room_name or len(room_name) > 16:
219
			raise ValidationError('Incorrect room name "{}"'.format(room_name))
220
		room = Room(name=room_name)
221
		do_db(room.save)
222
		RoomUsers(room_id=room.id, user_id=self.user_id).save()
223
		subscribe_message = self.subscribe_room_channel_message(room.id, room_name)
224
		self.publish(subscribe_message, self.channel, True)
225
226
	def invite_user(self, message):
227
		room_id = message[VarNames.ROOM_ID]
228
		user_id = message[VarNames.USER_ID]
229
		room = get_or_create_room(self.channels, room_id, user_id)
230
		users_in_room = {
231
			user.id: RedisPrefix.set_js_user_structure(user.username, user.sex)
232
			for user in room.users.all()
233
		}
234
		self.publish(self.add_user_to_room(room_id, user_id, users_in_room[user_id]), room_id)
235
		subscribe_message = self.invite_room_channel_message(room_id, user_id, room.name, users_in_room)
236
		self.publish(subscribe_message, RedisPrefix.generate_user(user_id), True)
237
238
	def create_room(self, user_rooms, user_id):
239
		if self.user_id == user_id:
240
			room_ids = list([room['room_id'] for room in self.evaluate(user_rooms)])
241
			query_res = execute_query(SELECT_SELF_ROOM, [room_ids, ])
242
		else:
243
			rooms_query = RoomUsers.objects.filter(user_id=user_id, room__in=user_rooms)
244
			query_res = rooms_query.values('room__id', 'room__disabled')
245
		try:
246
			room = do_db(query_res.get)
247
			room_id = room['room__id']
248
			update_room(room_id, room['room__disabled'])
249
		except RoomUsers.DoesNotExist:
250
			room_id = create_room_users(self.user_id, user_id)
251
		return room_id
252
253
	def create_user_channel(self, message):
254
		user_id = message[VarNames.USER_ID]
255
		# get all self private rooms ids
256
		user_rooms = Room.users.through.objects.filter(user_id=self.user_id, room__name__isnull=True).values('room_id')
257
		# get private room that contains another user from rooms above
258
		room_id = self.create_room(user_rooms, user_id)
259
		subscribe_message = self.subscribe_direct_channel_message(room_id, user_id)
260
		self.publish(subscribe_message, self.channel, True)
261
		other_channel = RedisPrefix.generate_user(user_id)
262
		if self.channel != other_channel:
263
			self.publish(subscribe_message, other_channel, True)
264
265
	def delete_channel(self, message):
266
		room_id = message[VarNames.ROOM_ID]
267
		if room_id not in self.channels or room_id == ALL_ROOM_ID:
268
			raise ValidationError('You are not allowed to exit this room')
269
		room = do_db(Room.objects.get, id=room_id)
270
		if room.disabled:
271
			raise ValidationError('Room is already deleted')
272
		if room.name is None:  # if private then disable
273
			room.disabled = True
274
		else:  # if public -> leave the room, delete the link
275
			RoomUsers.objects.filter(room_id=room.id, user_id=self.user_id).delete()
276
			online = self.get_online_from_redis(room_id)
277
			online.remove(self.user_id)
278
			self.publish(self.room_online(online, Actions.LOGOUT, room_id), room_id)
279
		room.save()
280
		message = self.unsubscribe_direct_message(room_id)
281
		self.publish(message, room_id, True)
282
283
	def edit_message(self, data):
284
		# ord(next (iter (message['images'])))
285
		message_id = data[VarNames.MESSAGE_ID]
286
		message = Message.objects.get(id=message_id)
287
		validate_edit_message(self.user_id, message)
288
		message.content = data[VarNames.CONTENT]
289
		selector = Message.objects.filter(id=message_id)
290
		if message.content is None:
291
			action = Actions.DELETE_MESSAGE
292
			prep_imgs = None
293
			selector.update(deleted=True)
294
		else:
295
			action = Actions.EDIT_MESSAGE
296
			prep_imgs = process_images(data.get(VarNames.IMG), message)
297
			selector.update(content=message.content, symbol=message.symbol)
298
		self.publish(self.create_send_message(message, action, prep_imgs), message.room_id)
299
300
	def send_client_new_channel(self, message):
301
		room_id = message[VarNames.ROOM_ID]
302
		self.add_channel(room_id)
303
		self.add_online_user(room_id)
304
305
	def send_client_delete_channel(self, message):
306
		room_id = message[VarNames.ROOM_ID]
307
		self.async_redis.unsubscribe((room_id,))
308
		self.async_redis_publisher.hdel(room_id, self.id)
309
		self.channels.remove(room_id)
310
311
	def process_get_messages(self, data):
312
		"""
313
		:type data: dict
314
		"""
315
		header_id = data.get(VarNames.GET_MESSAGES_HEADER_ID, None)
316
		count = int(data.get(VarNames.GET_MESSAGES_COUNT, 10))
317
		room_id = data[VarNames.CHANNEL]
318
		self.logger.info('!! Fetching %d messages starting from %s', count, header_id)
319
		if header_id is None:
320
			messages = Message.objects.filter(Q(room_id=room_id), Q(deleted=False)).order_by('-pk')[:count]
321
		else:
322
			messages = Message.objects.filter(Q(id__lt=header_id), Q(room_id=room_id), Q(deleted=False)).order_by('-pk')[:count]
323
		images = do_db(get_message_images, messages)
324
		response = self.get_messages(messages, room_id, images)
325
		self.ws_write(response)
326
327
328
class WebRtcMessageHandler(MessagesHandler):
329
330
	def __init__(self, *args, **kwargs):
331
		super(WebRtcMessageHandler, self).__init__(*args, **kwargs)
332
		self.pre_process_message.update({
333
			Actions.WEBRTC: self.proxy_webrtc,
334
			Actions.CLOSE_FILE_CONNECTION: self.close_file_connection,
335
			Actions.CLOSE_CALL_CONNECTION: self.close_call_connection,
336
			Actions.CANCEL_CALL_CONNECTION: self.cancel_call_connection,
337
			Actions.ACCEPT_CALL: self.accept_call,
338
			Actions.ACCEPT_FILE: self.accept_file,
339
			Actions.OFFER_FILE_CONNECTION: self.offer_webrtc_connection,
340
			Actions.OFFER_CALL_CONNECTION: self.offer_webrtc_connection,
341
			Actions.REPLY_FILE_CONNECTION: self.reply_file_connection,
342
			Actions.REPLY_CALL_CONNECTION: self.reply_call_connection,
343
		})
344
		self.post_process_message.update({
345
			Actions.OFFER_FILE_CONNECTION: self.set_opponent_call_channel,
346
			Actions.OFFER_CALL_CONNECTION: self.set_opponent_call_channel
347
		})
348
349
	def set_opponent_call_channel(self, message):
350
		connection_id = message[VarNames.CONNECTION_ID]
351
		if message[VarNames.WEBRTC_OPPONENT_ID] == self.id:
352
			return True
353
		self.sync_redis.hset(connection_id, self.id, WebRtcRedisStates.OFFERED)
354
355
	def offer_webrtc_connection(self, in_message):
356
		room_id = in_message[VarNames.CHANNEL]
357
		content = in_message.get(VarNames.CONTENT)
358
		qued_id = in_message[VarNames.WEBRTC_QUED_ID]
359
		connection_id = id_generator(RedisPrefix.CONNECTION_ID_LENGTH)
360
		# use list because sets dont have 1st element which is offerer
361
		self.async_redis_publisher.hset(WEBRTC_CONNECTION, connection_id, self.id)
362
		self.async_redis_publisher.hset(connection_id, self.id, WebRtcRedisStates.READY)
363
		opponents_message = self.offer_webrtc(content, connection_id, room_id, in_message[VarNames.EVENT])
364
		self_message = self.set_connection_id(qued_id, connection_id)
365
		self.ws_write(self_message)
366
		self.logger.info('!! Offering a webrtc, connection_id %s', connection_id)
367
		self.publish(opponents_message, room_id, True)
368
369
	def reply_file_connection(self, in_message):
370
		connection_id = in_message[VarNames.CONNECTION_ID]
371
		sender_ws_id = self.sync_redis.shget(WEBRTC_CONNECTION, connection_id)
372
		sender_ws_status = self.sync_redis.shget(connection_id, sender_ws_id)
373
		self_ws_status = self.sync_redis.shget(connection_id, self.id)
374
		if sender_ws_status == WebRtcRedisStates.READY and self_ws_status == WebRtcRedisStates.OFFERED:
375
			self.async_redis_publisher.hset(connection_id, self.id, WebRtcRedisStates.RESPONDED)
376
			self.publish(self.reply_webrtc(
377
				Actions.REPLY_FILE_CONNECTION,
378
				connection_id,
379
				HandlerNames.WEBRTC_TRANSFER
380
			), sender_ws_id)
381
		else:
382
			raise ValidationError("Invalid channel status.")
383
384
	def reply_call_connection(self, in_message):
385
		self.send_call_answer(
386
			in_message,
387
			WebRtcRedisStates.RESPONDED,
388
			Actions.REPLY_CALL_CONNECTION,
389
			[WebRtcRedisStates.OFFERED],
390
			HandlerNames.WEBRTC_TRANSFER
391
		)
392
393
	def proxy_webrtc(self, in_message):
394
		"""
395
		:type in_message: dict
396
		"""
397
		connection_id = in_message[VarNames.CONNECTION_ID]
398
		channel = in_message.get(VarNames.WEBRTC_OPPONENT_ID)
399
		self_channel_status = self.sync_redis.shget(connection_id, self.id)
400
		opponent_channel_status = self.sync_redis.shget(connection_id, channel)
401
		if not (self_channel_status == WebRtcRedisStates.READY and opponent_channel_status == WebRtcRedisStates.READY):
402
			raise ValidationError('Error in connection status, your status is {} while opponent is {}'.format(
403
				self_channel_status, opponent_channel_status
404
			))  # todo receiver should only accept proxy_webrtc from sender, sender can accept all
405
		# I mean somebody if there're 3 ppl in 1 channel and first is initing transfer to 2nd and 3rd,
406
		# 2nd guy can fraud 3rd guy webrtc traffic, which is allowed during the call, but not while transering file
407
		in_message[VarNames.WEBRTC_OPPONENT_ID] = self.id
408
		in_message[VarNames.HANDLER_NAME] = HandlerNames.PEER_CONNECTION
409
		self.logger.debug(
410
			"Forwarding message to channel %s, self %s, other status %s",
411
			channel,
412
			self_channel_status,
413
			opponent_channel_status
414
		)
415
		self.publish(in_message, channel)
416
417
	def close_file_connection(self, in_message):
418
		connection_id = in_message[VarNames.CONNECTION_ID]
419
		self_channel_status = self.sync_redis.shget(connection_id, self.id)
420
		if not self_channel_status:
421
			raise Exception("Access Denied")
422
		if self_channel_status != WebRtcRedisStates.CLOSED:
423
			sender_id = self.sync_redis.shget(WEBRTC_CONNECTION, connection_id)
424
			if sender_id == self.id:
425
				self.close_file_sender(connection_id)
426
			else:
427
				self.close_file_receiver(connection_id, in_message, sender_id)
428
			self.async_redis_publisher.hset(connection_id, self.id, WebRtcRedisStates.CLOSED)
429
430
	def close_call_connection(self, in_message):
431
		self.send_call_answer(
432
			in_message,
433
			WebRtcRedisStates.CLOSED,
434
			Actions.CLOSE_CALL_CONNECTION,
435
			[WebRtcRedisStates.READY, WebRtcRedisStates.RESPONDED],
436
			HandlerNames.PEER_CONNECTION
437
		)
438
439
	def cancel_call_connection(self, in_message):
440
		self.send_call_answer(
441
			in_message,
442
			WebRtcRedisStates.CLOSED,
443
			Actions.CANCEL_CALL_CONNECTION,
444
			[WebRtcRedisStates.OFFERED],
445
			HandlerNames.WEBRTC_TRANSFER
446
		)
447
448
	def close_file_receiver(self, connection_id, in_message, sender_id):
449
		sender_status = self.sync_redis.shget(connection_id, sender_id)
450
		if not sender_status:
451
			raise Exception("Access denied")
452
		if sender_status != WebRtcRedisStates.CLOSED:
453
			in_message[VarNames.WEBRTC_OPPONENT_ID] = self.id
454
			in_message[VarNames.HANDLER_NAME] = HandlerNames.PEER_CONNECTION
455
			self.publish(in_message, sender_id)
456
457
	def close_file_sender(self, connection_id):
458
		values = self.sync_redis.shgetall(connection_id)
459
		del values[self.id]
460
		message = self.get_close_file_sender_message(connection_id)
461
		for ws_id in values:
462
			if values[ws_id] == WebRtcRedisStates.CLOSED:
463
				continue
464
			self.publish(message, ws_id)
465
466
	def accept_file(self, in_message):
467
		connection_id = in_message[VarNames.CONNECTION_ID]
468
		sender_ws_id = self.sync_redis.shget(WEBRTC_CONNECTION, connection_id)
469
		sender_ws_status = self.sync_redis.shget(connection_id, sender_ws_id)
470
		self_ws_status = self.sync_redis.shget(connection_id, self.id)
471
		if sender_ws_status == WebRtcRedisStates.READY and self_ws_status == WebRtcRedisStates.RESPONDED:
472
			self.async_redis_publisher.hset(connection_id, self.id, WebRtcRedisStates.READY)
473
			self.publish(self.get_accept_file_message(connection_id), sender_ws_id)
474
		else:
475
			raise ValidationError("Invalid channel status")
476
477
	# todo
478
	# we can use channel_status = self.sync_redis.shgetall(connection_id)
479
	# and then self.async_redis_publisher.hset(connection_id, self.id, WebRtcRedisStates.READY)
480
	# if we shgetall and only then do async hset
481
	# we can catch an issue when 2 concurrent users accepted the call
482
	# but we didn't  send them ACCEPT_CALL as they both were in status 'offered'
483
	def accept_call(self, in_message):
484
		connection_id = in_message[VarNames.CONNECTION_ID]
485
		self_status = self.sync_redis.shget(connection_id, self.id)
486
		if self_status == WebRtcRedisStates.RESPONDED:
487
			conn_users = self.sync_redis.shgetall(connection_id)
488
			self.publish_call_answer(
489
				conn_users,
490
				connection_id,
491
				HandlerNames.WEBRTC_TRANSFER,
492
				Actions.ACCEPT_CALL,
493
				WebRtcRedisStates.READY
494
			)
495
		else:
496
			raise ValidationError("Invalid channel status")
497
498
	def send_call_answer(self, in_message, status_set, reply_action, allowed_state, message_handler):
499
		connection_id = in_message[VarNames.CONNECTION_ID]
500
		conn_users = self.sync_redis.shgetall(connection_id)
501
		if conn_users[self.id] in allowed_state:
502
			self.publish_call_answer(conn_users, connection_id, message_handler, reply_action, status_set)
503
		else:
504
			raise ValidationError("Invalid channel status.")
505
506
	def publish_call_answer(self, conn_users, connection_id, message_handler, reply_action, status_set):
507
		self.async_redis_publisher.hset(connection_id, self.id, status_set)
508
		del conn_users[self.id]
509
		message = self.reply_webrtc(reply_action, connection_id, message_handler)
510
		for user in conn_users:
511
			if conn_users[user] != WebRtcRedisStates.CLOSED:
512
				self.publish(message, user)