Client.py

import time, datetime, ip2country
import logging


# Spring and Recoil players change unit stats for one battle by sending SPADS a compiled
# tweak set as chat, split over numbered slots: "!bset tweakdefs<n> <base64>". A slot is
# 16000 base64 characters in the tooling players use today, far over any account's normal
# message length, so those two commands get their own limit. Teiserver allows 16385 for the
# same payload. The number here also has to cover the "SAYBATTLE " prefix, which teiserver's
# equivalent check has already stripped.
TWEAK_MSG_LENGTH = 16384
TWEAK_COMMAND_PREFIXES = ('!bset tweakdefs', '!bset tweakunits')

def is_tweak_command(command):
	'whether one raw command line is a SPADS tweak set going to a battle'
	# the line is still raw here: it may carry a "#<id> " message id, and the command
	# name is whatever case the client sent it in
	if command.startswith('#'):
		msg_id, _, rest = command.partition(' ')
		if msg_id[1:].isdigit():
			command = rest

	name, _, msg = command.partition(' ')
	if name.upper() not in ('SAYBATTLE', 'SAYBATTLEEX'):
		return False
	return msg.lower().startswith(TWEAK_COMMAND_PREFIXES)

def command_length_limit(command, flood_limits, logged_in):
	'the message length limit that applies to one raw command line'
	limit = flood_limits['msglength']
	if not logged_in or not is_tweak_command(command):
		return limit
	return max(limit, TWEAK_MSG_LENGTH)

def tweak_command_bytes(lines, logged_in):
	'''
	How many of a chunk's bytes belong to tweak commands, and so are charged to the
	tweak allowance rather than to the byte rate.

	A line over the tweak length limit is dropped later on and never reaches a battle, so
	it earns no exemption. Without that, a flood only has to wear a tweak command's prefix
	to escape the byte rate entirely.
	'''
	if not logged_in:
		return 0
	# each line was charged with the newline that ended it
	return sum(len(line) + 1 for line in lines
		if len(line) <= TWEAK_MSG_LENGTH and is_tweak_command(line))


class Client():
	'this object represents one server-side connected client'

	def __init__(self, root, address, session_id):
		'initial setup for the connected client'
		self._root = root
		now = time.time()

		# detects if the connection is from this computer
		if address[0].startswith('127.'):
			if root.online_ip:
				address = (root.online_ip, address[1])
			elif root.local_ip:
				address = (root.local_ip, address[1])

		self.ip_address = address[0]
		self.local_ip = address[0]
		self.port = address[1]

		# fields also in user db
		self.user_id = -1 # db user object has a .id attr instead
		self.username = ""
		self.password = ""
		self.register_date = datetime.datetime.now()
		self.last_login = datetime.datetime.now()
		self.last_ip = self.ip_address
		self.last_id = 0
		self.ingame_time = 0
		self.access = 'fresh'
		self.email = ''
		self.bot = False

		# session
		self.session_id = session_id
		self.debug = False
		self.static = False
		self.sendError = False

		self.compat = set() # holds compatibility flags

		self.country_code = '??'
		self.agent = ""
		self.setFlagByIP(self.ip_address)
		self.status = 12
		self.accesslevels = ['fresh','everyone']

		# note: this NEVER becomes false after LOGIN!
		self.logged_in = False

		# server<->client comms
		self.buffersend = False # if True, write all sends to a buffer (must not be used when a client is logging in but didn't yet receive full server state!)
		self.buffer = [] # list of message strings (incl. trailing newline); joined once in flushBuffer
		self.msg_id = ''
		self.msg_sendbuffer = []
		self.sendingmessage = ''
		self.msg_length_history = {}
		self.tweak_length_history = {} # the part of it spent on tweak sets, exempt from the byte rate

		# channels
		self.channels = set()
		self.ignored = {}
		self.lastsaid = {}
		
		# for if we are a bridge bot
		self.bridge = {} #location->{external_id->bridged_id}
		
		# perhaps these are unused?
		self.cpu = 0
		self.data = ''
		self.lastdata = now

		# time-stamps for encrypted data
		self.incoming_msg_ctr = 0
		self.outgoing_msg_ctr = 1

		# battle stuff
		self.is_ingame = False
		self.scriptPassword = None

		self.battle_bots = {}
		self.current_battle = None # battle_id
		self.pending_battle = None # battle_id
		self.relayed_host_ip = None # address from RELAYEDHOST, waiting for the next OPENBATTLE
		self.went_ingame = 0
		self.spectator = False
		self.battlestatus = {'ready':'0', 'id':'0000', 'ally':'0000', 'mode':'0', 'sync':'00', 'side':'00', 'handicap':'0000000'}
		self.teamcolor = '0'

		self.hostport = None
		self.udpport = 0

	def set_msg_id(self, msg):
		self.msg_id = ""

		if (not msg.startswith('#')):
			return msg

		test = msg.split(' ')[0][1:]

		if (not test.isdigit()):
			return msg

		self.msg_id = '#%s ' % test
		return (' '.join(msg.split(' ')[1:]))

	def setFlagByIP(self, ip, force=True):
		cc = ip2country.lookup(ip)
		if force or cc != '??':
			self.country_code = cc

	##
	## handle data from client
	##
	def Handle(self, data):
		if self.bot:
			flood_limits = self._root.flood_limits['bot']
		elif (self.access in self._root.flood_limits):
			flood_limits = self._root.flood_limits[self.access]
		else:
			flood_limits = self._root.flood_limits['fresh']
		
		#logging.info("  < [" + self.username + " " + str(self.session_id) + "] " + data.strip()) # uncomment for debugging

		now = int(time.time())
		self.lastdata = now # data received, store time to detect disconnects

		bytespersecond = flood_limits['bytespersecond']
		seconds = flood_limits['seconds']

		# keep appending until we see at least one newline
		self.data += data
		split_data = self.data.split("\n")

		if (now in self.msg_length_history):
			self.msg_length_history[now] += len(data)
		else:
			self.msg_length_history[now] = len(data)

		total = self.sumFloodHistory(self.msg_length_history, now, seconds)
		exempt = self.sumFloodHistory(self.tweak_length_history, now, seconds)

		# A tweak set is several 16k slots sent back to back, which is over the byte rate of
		# every account type, so a player sending one would be disconnected partway through.
		# Those bytes are charged to a separate allowance instead. Once it is spent the rest
		# is charged as ordinary traffic, so a flood wearing a tweak command's prefix still
		# meets the same limit it always did. Only the lines this call completed are counted,
		# and each of them only once.
		spare = max(0, flood_limits['tweakbytes'] - exempt)
		tweaked = min(spare, tweak_command_bytes(split_data[: len(split_data) - 1], self.logged_in))
		self.tweak_length_history[now] = self.tweak_length_history.get(now, 0) + tweaked
		exempt += tweaked

		if (total - exempt) > (bytespersecond * seconds):
			self.Send('SERVERMSG No flooding (over %s per second for %s seconds)' % (bytespersecond, seconds))
			self.ReportFloodBreach("flood limit", total - exempt)
			self.Remove('Kicked for flooding (%s)' % (self.access))
			return

		# if far too much data has accumulated without hitting flood limits and without a newline, just clear it
		if (len(split_data) == 1):
			if (len(self.data) > (flood_limits['msglength'])*16):
				del self.data
				self.data = ""
				self.Send('SERVERMSG Max client data cache was exceeded, some of your data was dropped by the server')
				self.ReportFloodBreach("max client data cache ", len(self.data))
			return

		self.HandleProtocolCommands(split_data, flood_limits)

	def sumFloodHistory(self, history, now, seconds):
		'total of one per-second byte history over the flood window, dropping what fell out of it'
		total = 0
		for iter in dict(history):
			if (iter < now - (seconds - 1)):
				del history[iter]
			else:
				total += history[iter]
		return total

	def HandleProtocolCommand(self, cmd):
		# probably caused by trailing newline ("abc\n".split("\n") == ["abc", ""])
		if (len(cmd) < 1):
			return
		self._root.protocol._handle(self, cmd)

	def HandleProtocolCommands(self, split_data, flood_limits):
		assert(type(split_data) == list)
		assert(type(split_data[-1]) == str)

		# either a list of commands, or a list of encrypted data
		# blobs which may contain embedded (post-decryption) NLs
		# note: will be empty if len(split_data) == 1
		raw_data_blobs = split_data[: len(split_data) - 1]

		# will be a single newline in most cases, or an incomplete
		# command which should be saved for a later time when more
		# data is in buffer
		self.data = split_data[-1]

		commands_buffer = []

		for raw_data_blob in raw_data_blobs:
			if (len(raw_data_blob) == 0):
				continue

			strip_commands = [(raw_data_blob.rstrip('\r')).lstrip(' ')]
			commands_buffer += strip_commands

		for command in commands_buffer:
			length_limit = command_length_limit(command, flood_limits, self.logged_in)
			if len(command) > length_limit:
				self.Send('SERVERMSG message length limit of %i chars was exceeded: command \"%s...\" dropped.' % (length_limit, command[0: 16]))
				self.ReportFloodBreach("max message length (cmd=%s...)" % command[0: 16], len(command))
				continue
			self.HandleProtocolCommand(command)

	def ReportFloodBreach(self, type, bytes):
		if hasattr(self, "username"):
			user_details = "<%s>, session_id: %i" % (self.username, self.session_id)
		else:
			user_details = "session_id: %i" % self.session_id
		err_msg = "%s for '%s' breached by %s, had %i bytes" % (type, self.access, user_details, bytes)
		self._root.protocol.broadcast_Moderator(err_msg)
		logging.info(err_msg)

	##
	## send data to client
	##
	def RealSend(self, data, command=None):
		if not data:
			return

		# 2.1: command may be precomputed by the broadcast path (same token for every
		# recipient of one message). Only fall back to deriving it per-send otherwise.
		if command is None:
			raw_msg  = data[data.find(" ")+1:] if data.startswith('#') else data
			command = raw_msg[:raw_msg.find(" ")] if " " in raw_msg else raw_msg
		self._root.outbound_command_stats[command] = self._root.outbound_command_stats.get(command, 0) + 1

		#logging.info("> [" + self.username + " " + str(self.session_id) + "] " + data.strip()) # uncomment for debugging

		# while buffersend is on (the login state-dump), accumulate into the buffer
		# and let flushBuffer emit it in a single transport.write. Stats are counted
		# above either way, and no msg_id is prepended, so the bytes are identical to
		# sending each message directly - only the syscall batching differs.
		if self.buffersend:
			self.buffer.append(data + "\n")
		else:
			self.transport.write(data.encode("utf-8") + b"\n")

	def Send(self, data, command=None):
		if self.msg_id:
			data = self.msg_id + data
		if self.buffersend:
			self.buffer.append(data + "\n")
		else:
			self.RealSend(data, command)

	def flushBuffer(self):
		self.transport.write("".join(self.buffer).encode("utf-8"))
		self.buffer = []
		self.buffersend = False

	def isAdmin(self):
		return ('admin' in self.accesslevels)

	def isMod(self):
		return self.isAdmin() or ('mod' in self.accesslevels) # maybe cache these
		
	def isHosting(self):
		return self.current_battle and self._root.battles[self.current_battle].host == self.session_id


def selftest():
	limits = {'msglength': 10000}
	payload = 'A' * 16000

	def limit_for(cmd, logged_in = True):
		return command_length_limit(cmd, limits, logged_in)

	# tweak commands get the raised limit, on both say paths and with a message id
	assert(limit_for('SAYBATTLE !bset tweakunits1 ' + payload) == TWEAK_MSG_LENGTH)
	assert(limit_for('SAYBATTLEEX !bset tweakdefs10 ' + payload) == TWEAK_MSG_LENGTH)
	assert(limit_for('#42 SAYBATTLE !bset tweakdefs1 ' + payload) == TWEAK_MSG_LENGTH)
	assert(limit_for('saybattle !bset TweakUnits1 ' + payload) == TWEAK_MSG_LENGTH)

	# everything else stays on the account limit
	assert(limit_for('SAYBATTLE hello') == limits['msglength'])
	assert(limit_for('SAY #main !bset tweakdefs1 ' + payload) == limits['msglength'])
	assert(limit_for('SAYPRIVATE host !bset tweakdefs1 ' + payload) == limits['msglength'])
	assert(limit_for('SAYBATTLE !bset tweakdefs1 ' + payload, False) == limits['msglength'])

	# the raise never lowers an account limit that is already higher
	assert(command_length_limit('SAYBATTLE !bset tweakdefs1', {'msglength': 99999}, True) == 99999)

	# and the length check in HandleProtocolCommands uses it
	class FakeClient(Client):
		def __init__(self):
			self.data = ''
			self.logged_in = True
			self.handled = []
			self.sent = []
		def HandleProtocolCommand(self, cmd):
			self.handled.append(cmd)
		def Send(self, data, command = None):
			self.sent.append(data)
		def ReportFloodBreach(self, type, bytes):
			pass

	tweak = FakeClient()
	tweak.HandleProtocolCommands(['SAYBATTLE !bset tweakunits1 ' + payload, ''], limits)
	assert(len(tweak.handled) == 1)
	assert(tweak.sent == [])

	chat = FakeClient()
	chat.HandleProtocolCommands(['SAYBATTLE ' + payload, ''], limits)
	assert(chat.handled == [])
	assert(len(chat.sent) == 1)

	# a whole tweak set is several slots sent back to back, far over a player's byte rate
	rate = {'msglength': 10000, 'bytespersecond': 2000, 'seconds': 10,
		'tweakbytes': 60 * TWEAK_MSG_LENGTH}

	class FloodClient(Client):
		def __init__(self, flood_limits, logged_in = True):
			class FakeRoot:
				pass
			self._root = FakeRoot()
			self._root.flood_limits = {'user': flood_limits, 'fresh': flood_limits}
			self.bot = False
			self.access = 'user'
			self.logged_in = logged_in
			self.data = ''
			self.msg_length_history = {}
			self.tweak_length_history = {}
			self.handled = []
			self.sent = []
			self.removed = []
		def HandleProtocolCommand(self, cmd):
			self.handled.append(cmd)
		def Send(self, data, command = None):
			self.sent.append(data)
		def ReportFloodBreach(self, type, bytes):
			pass
		def Remove(self, reason = 'Quit'):
			self.removed.append(reason)

	def slots(n):
		return ["SAYBATTLE !bset tweakunits%d %s\n" % (i, payload) for i in range(1, n + 1)]

	# five slots, one write each: all delivered, sender still connected
	paced = FloodClient(rate)
	for line in slots(5):
		paced.Handle(line)
	assert(len(paced.handled) == 5)
	assert(paced.removed == [])

	# and the same five arriving in one read, which is what a client blasting them looks like
	coalesced = FloodClient(rate)
	coalesced.Handle("".join(slots(5)))
	assert(len(coalesced.handled) == 5)
	assert(coalesced.removed == [])

	# the byte rate is untouched for everything else: the same volume of ordinary chat
	# still disconnects the sender
	flooder = FloodClient(rate)
	for _ in range(5):
		flooder.Handle("SAYBATTLE %s\n" % payload)
	assert(flooder.removed != [])

	# past the allowance, tweak bytes are charged like any others
	stingy = FloodClient(dict(rate, tweakbytes = 2 * TWEAK_MSG_LENGTH))
	for line in slots(5):
		stingy.Handle(line)
	assert(stingy.removed != [])

	# a line too long to be delivered anyway buys no exemption
	overlong = FloodClient(rate)
	for _ in range(5):
		overlong.Handle("SAYBATTLE !bset tweakunits1 %s\n" % ('A' * TWEAK_MSG_LENGTH))
	assert(overlong.removed != [])

	# and neither does a tweak command from a client that has not logged in
	anon = FloodClient(rate, logged_in = False)
	for line in slots(5):
		anon.Handle(line)
	assert(anon.removed != [])

	print("Client.py selftest passed")

if __name__ == '__main__':
	selftest()