test_pbms_client.rb

#-----------------------------------------------------------------------------
# Test PBMS master client
# ( use with https://github.com/jamesu/PushButton-Master-Server )
#
# Copyright (C) 2016 James S Urquhart.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#-----------------------------------------------------------------------------

require 'socket'
require 'stringio'
require 'ipaddr'

$sock = UDPSocket.new(Socket::AF_INET6)
#$sock.bind("::1", 28002)
$sock.bind("::1", 28003)


class Packet
	Type_Normal = 0
	Type_Buddy= 1
	Type_Offline = 2
	Type_Favourites = 3

	QueryFlags_OnlineQuery = 0
	QueryFlags_OfflineQuery = 0x1
	QueryFlags_NoStringCompress = 0x2
	QueryFlags_NewStyleResponse = 0x4 # implies extended header
	QueryFlags_Authenticated = 0x8

	FilterFlags_Dedicated = 0x1
	FilterFlags_NotPassworded = 0x2
	FilterFlags_Linux = 0x4
	FilterFlags_CurrentVersion = 0x8
	FilterFlags_NotXenon = 0x10

	RegionMask_RegionIsIPV4Address = 0x40000000
	RegionMask_RegionIsIPV6Address = 0x80000000

	ServerAddress_IPV4 = 0
	ServerAddress_IPV6 = 1


	def self.implement_packet(ident)
		@@packet_handlers ||= {}
		@@packet_handlers[ident] = self
	end

	def self.packet_handlers
		@@packet_handlers
	end

	def self.recv_identify(data)
		type = data.unpack('C')
	end

	def self.on_packet(type, &block)
		@@packet_proc_handlers ||= {}
		@@packet_proc_handlers[type] = block
	end

	def self.handle_packet(inst)
		if @@packet_proc_handlers.has_key?(inst.class)
			@@packet_proc_handlers[inst.class].call(inst)
		end
	end

	def writeCString(stream, str)
		str = str.encode('utf-8')
		stream.write([str.length].pack('C'))
		stream.write(str)
	end

	def to_string
		sio = StringIO.new
		write(sio)
		return sio.string
	end

	def readHeader(stream)
		typeid, flags = stream.read(2).unpack('CC')
		if (flags & Packet::QueryFlags_Authenticated) != 0
			puts "READING AUTHENTICATED HEADER"
			session = stream.read(4).unpack('L')[0]
			return [typeid, flags, session, 0]
		else
			puts "READING NORMAL HEADER"
			session, key = stream.read(4).unpack('SS')
			return [typeid, flags, session, key]
		end
	end

	def writeHeader(stream, typeid, in_flags, session, key)
		if (in_flags & Packet::QueryFlags_Authenticated) != 0
			puts "WRITING AUTHENTICATED HEADER #{flags}"
			stream.write([typeid, in_flags, session].pack('CCL'))
		else
			puts "WRITING NORMAL HEADER"
			stream.write([typeid, in_flags, session, key].pack('CCSS'))
		end
	end
end

class MasterServerListRequest < Packet
	implement_packet 6

	attr_accessor :flags
	attr_accessor :session
	attr_accessor :key
	attr_accessor :gameType
	attr_accessor :missionType
	attr_accessor :minPlayers
	attr_accessor :maxPlayers
	attr_accessor :regionMask
	attr_accessor :version
	attr_accessor :maxBots
	attr_accessor :minCPU
	attr_accessor :buddies
	attr_accessor :index

	def initialize
		@flags = 0
		@session = 0
		@key = 0
		@gameType = 't3d'
		@missionType = 'normal'
		@minPlayers = 0
		@maxPlayers = 0
		@regionMask = 0
		@version = 0
		@filterFlags = 0
		@maxBots = 0
		@minCPU = 0
		@buddies = []
		@index = nil
	end

	def read(stream)
		[]
	end

	def write(stream)

		if index != nil
			puts "INDEX IS #{@index}"
			writeHeader(stream, 6, @flags, @session, @key)

			stream.write([
				@index,
				0,
				0,
				0,
				0,
				0,
				0,
				0,
				0,
				0,
				0
				].pack("CCCLLCCSC"))
			return
		end

		puts "INDEX IS NIL"
		writeHeader(stream, 6, @flags, @session, @key)

		stream.write([
			255,
			].pack("C"))

		writeCString(stream, @gameType)
		writeCString(stream, @missionType)

		stream.write([
			@minPlayers,
			@maxPlayers,
			@regionMask,
			@version,
			@filterFlags,
			@maxBots,
			@minCPU,
			@buddies.length].pack("CCLLCCSC"))
		
		if @buddies.length > 0
			stream.write(@buddies.pack("L*"))
		end
	end
end

class MasterServerExtendedListRequest < MasterServerListRequest
	implement_packet 44

	def write(stream)

		if index != nil
			puts "INDEX IS #{@index}"
			writeHeader(stream, 44, @flags, @session, @key)
				stream.write([
					@index,
					0,
					0,
					0,
					0,
					0,
					0,
					0,
					0,
					0,
					0
					].pack("SCCLLCCSC"))

			return
		end

		puts "INDEX IS NIL"
		writeHeader(stream, 44, @flags, @session, @key)

		stream.write([
			65535,
			].pack("S"))

		writeCString(stream, @gameType)
		writeCString(stream, @missionType)

		stream.write([
			@minPlayers,
			@maxPlayers,
			@regionMask,
			@version,
			@filterFlags,
			@maxBots,
			@minCPU,
			@buddies.length].pack("CCLLCCSC"))
		
		if @buddies.length > 0
			stream.write(@buddies.pack("L*"))
		end
	end
end

class MasterServerListResponse < Packet
	implement_packet 8

	attr_accessor :flags
	attr_accessor :session
	attr_accessor :key
	attr_accessor :index
	attr_accessor :totalPackets
	attr_accessor :servers

	def initialize
		@flags = 0
		@session = 0
		@key = 0
		@index = 0
		@totalPackets = 0
		@servers = []
	end

	def read(stream)
		# header == 6 bytes (type -> key) CCSS
		# 
		type, @flags, @session, @key = readHeader(stream)
		@index, @totalPackets, serverCount = stream.read(4).unpack('CCS')
		puts "#{serverCount} servers"
		@servers = []
		serverCount.times {
			data = stream.read(6).unpack('CCCCS')
			@servers << [Packet::ServerAddress_IPV4, data[0..3], data[4]]
		}
	end
end

class MasterServerExtendedListResponse < Packet
	implement_packet 40

	attr_accessor :flags
	attr_accessor :session
	attr_accessor :key
	attr_accessor :index
	attr_accessor :totalPackets
	attr_accessor :servers

	def initialize
		@flags = 0
		@session = 0
		@key = 0
		@index = 0
		@totalPackets = 0
		@servers = []
	end

	def read(stream)
		type, @flags, @session, @key = readHeader(stream)
		if (@flags & Packet::QueryFlags_NewStyleResponse) == 0
			puts "WTF WTF EXTENDED PACKET WITH NO NEWSTYLERESPONSE SET!!"
			exit
		end

		@index, @totalPackets, serverCount = stream.read(6).unpack('SSS')
		@servers = []
		puts "#{serverCount} servers index == #{@index}, totalPackets == #{@totalPackets}"
		serverCount.times {
			type = stream.read(1).unpack('C')[0]
			if type == 0
				data = stream.read(6).unpack('CCCCS')
				@servers << [Packet::ServerAddress_IPV4, data[0..3], data[4]]
			elsif type == 1
				data = stream.read(18).unpack('CCCCCCCCCCCCCCCCCS')
				@servers << [Packet::ServerAddress_IPV6, data[0..15], data[16]]
			else
				puts "Warning: unknown server type #{type}"
			end
		}
	end
end

class MasterServerChallengePacket < Packet
	implement_packet 42

	attr_accessor :flags
	attr_accessor :token
	attr_accessor :client_token
	attr_accessor :session
	attr_accessor :key

	def initialize
		@flags = 0
		@session = 0
		@key = 0
		@client_token = 0
	end

	def read(stream)
		type, @flags, @client_token, @session, @key = stream.read(10).unpack('CCLSS')
	end

	def write(stream)
		stream.write([42, @flags, @client_token, @session, @key].pack('CCLSS'))
	end

	# Requests a new challenge from the server
	def begin(session, key)
		@sess = session
		@key = key
		@flags = 0
	end

	# Returns a MasterServerChallengePacket which answers the challenge
	def get_client_session()
		return @client_token
	end
end

class PendingServerQueryResponse
	attr_accessor :start_time
	attr_accessor :retry_count
	attr_accessor :packet
	attr_accessor :typeId

	def initialize(typeId)
		@typeId = typeId
		@start_time = Time.now
		@retry_count = 0
		@packet = nil
	end
end

# Query timed out
class ServerQueryTimedOut < Exception
end

# Session prematurely ended, query needs to be restarted
class ServerQuerySessionEnded < Exception
end

# Server sent an unexpected packet
class ServerQueryBadResponse < Exception
end

class ServerQuery
	attr_accessor :session
	attr_accessor :key
	attr_accessor :extended_response
	attr_accessor :response

	# Challenge packet related stuff
	attr_accessor :use_challenge
	attr_accessor :authenticated
	attr_accessor :challenge_sent_packet
	attr_accessor :challenge_session

	TIMEOUT_SECONDS = 5
	MAX_PACKET_RETRY = 1

	def authenticated
		@authenticated
	end

	def authenticated_session
		@challenge_session.nil? ? @session : @challenge_session
	end

	def initialize(socket)
		@session = 0
		@key = 0
		@extended_response = false
		@response = {}
		@packet_count = 0
		@socket = socket

		@use_challenge = false
		@challenge_sent_packet = nil
		@challenge_session = nil
	end

	def send_challenge()
		puts "Sending challenge"
		@response[0] = PendingServerQueryResponse.new(:challenge)

		@challenge_sent_packet = MasterServerChallengePacket.new
		@challenge_sent_packet.begin(@session, @key)

		@socket.send(@challenge_sent_packet.to_string, 0, @server, @port)
	end

	def send_query()
		@response[0] = PendingServerQueryResponse.new(:list)

		if !@challenge_session.nil?
			@query.session = @challenge_session
			@query.flags |= Packet::QueryFlags_Authenticated | Packet::QueryFlags_NewStyleResponse
		else
			@query.session = @session
			@query.flags &= ~Packet::QueryFlags_Authenticated
		end

		puts "Sending query SESSION == #{@query.session}"
		@socket.send(@query.to_string, 0, @server, @port)
	end

	def query_servers(server, port, query)
		@session = query.session
		@key = query.key
		@extended_response = false
		@response = {}
		@packet_count = 0
		@server = server
		@port = port
		@query = query.clone
		@authenticated = !@use_challenge
		@challenge_sent_packet = nil
		@challenge_session = nil

		Packet.on_packet(MasterServerChallengePacket) do |packet|
			@authenticated = packet.flags & Packet::QueryFlags_Authenticated

			puts "PACKET SESSION #{packet.session} KEY #{packet.key} NEWKEY #{packet.client_token} :: #{@session} / #{@key}"
			if (packet.session == @session or packet.session == @challenge_session) && packet.key == @key
				if !@authenticated
					if @packet_count > 0
						# If we're in the middle of a list, we no longer have any session data
						throw ServerQuerySessionEnded.new
					end

					@challenge_session = nil
					puts "Received challenge, not authenticated. Trying again..."
					send_challenge()
				else
					puts "Challenge authenticated"
					@challenge_session = packet.get_client_session

					if @packet_count == 0
						puts "Continuing with query"
						send_query()
					else
						puts "Weird challenge case, shouldn't happen."
						throw ServerQuerySessionEnded.new
					end
				end
			else
				puts "Received bad challenge packet from master #{packet.session}/#{packet.key}"
				throw ServerQueryBadResponse.new
			end
		end

		Packet.on_packet(MasterServerListResponse) do |packet|
			@extended_response = false
			if packet.session == authenticated_session
				on_normal_response(packet)
			else
				on_erroneous_response(packet)
			end
		end

		Packet.on_packet(MasterServerExtendedListResponse) do |packet|
			if packet.session == authenticated_session
				@extended_response = true
				on_extended_response(packet)
			else
				on_erroneous_response(packet)
			end
		end

		# Start base request
		if @authenticated
			send_query()
		else
			send_challenge()
		end
	end

	def finished_query?
		@response.each do |k,v|
			if v.packet.nil?
				return false
			end
		end

		return true
	end

	def response_timed_out?
		now = Time.now

		@response.each do |k,v|
			if v.packet.nil?
				if ((now - v.start_time) > ServerQuery::TIMEOUT_SECONDS) && v.retry_count >= ServerQuery::MAX_PACKET_RETRY
					return true
				end
			end
		end

		return false
	end

	def issue_retry_packets
		now = Time.now
		
		@response.each do |k,v|
			if v.packet.nil?
				#puts "STILL WAITING FOR PACKET #{k}"
				if ((now - v.start_time) > ServerQuery::TIMEOUT_SECONDS) && v.retry_count < ServerQuery::MAX_PACKET_RETRY
					if v.typeId == :challenge
						puts "Retrying challenge"
						v.start_time = Time.now
						v.retry_count = v.retry_count + 1
						@socket.send(@challenge_sent_packet.to_string, 0, @server, @port)
					else
						puts "Retrying list packet #{k}"
						v.start_time = Time.now
						v.retry_count = v.retry_count + 1
						@query.index = k
						@socket.send(@query.to_string, 0, @server, @port)
					end
				end
			end
		end
	end

	def wait_for_full_response(&block)
		while !finished_query?
			block.call()

			if response_timed_out?
				throw ServerQueryTimedOut.new
			end

			issue_retry_packets
		end
	end

	def ensure_packets_are_reserved
		(0...@packet_count).each do |i|
			if @response[i].nil?
				puts "Reserved slot for packet #{i}"
				@response[i] = PendingServerQueryResponse.new(:list)
			end
		end
	end

	def on_normal_response(packet)
		@packet_count = packet.totalPackets
		ensure_packets_are_reserved

		if @response[packet.index].nil?
			puts "Warning: got unexpected packet #{packet.index}"
			return
		end

		puts "Setting packet #{packet.index}"
		@response[packet.index].packet = packet
	end

	def on_extended_response(packet)
		@packet_count = packet.totalPackets
		ensure_packets_are_reserved

		if @response[packet.index].nil?
			puts "Warning: got unexpected packet #{packet.index}"
			return
		end

		puts "Setting packet #{packet.index}"
		@response[packet.index].packet = packet
	end

	def on_erroneous_response(packet)
		puts "Got response packet with session #{packet.session} key #{packet.key}"
		exit
	end

	def servers
		out_list = []
		(0...@packet_count).each do |index|
			next if !@response.has_key?(index)
			next if @response[index].packet.nil?
			next if @response[index].typeId != :list
			out_list += @response[index].packet.servers
		end

		return out_list
	end
end


def handle_socket(socket)
	begin # emulate blocking recvfrom
	  msg, addr = socket.recvfrom_nonblock(1500)

	  packet_type = msg.unpack('C')[0]
	  #puts msg.inspect
	  #puts "PACKET TYPE IS #{packet_type}"

	  klass = Packet.packet_handlers[packet_type]
	  #puts klass.inspect
	  if !klass.nil?
	  	instance = klass.new
	  	File.open('last_packet.raw', 'wb') { |f| f.write(msg) }
	  	instance.read(StringIO.new(msg))
	  	Packet.handle_packet(instance)
	  end

	rescue IO::WaitReadable
	  #IO.select([socket])
	  #retry
	end
end

$query = ServerQuery.new($sock)
$query_data = MasterServerListRequest.new
#$query_data.queryFlags = Packet::QueryFlags_NewStyleResponse
$query_data.gameType = 'TEST'
$query_data.missionType = 'NORMAL'
#$query_data.regionMask = Packet::RegionMask_RegionIsIPV6Address | Packet::RegionMask_RegionIsIPV4Address
#$query_data.regionMask = Packet::RegionMask_RegionIsIPV4Address
#$query.use_challenge = true
$query.query_servers('::1', 28002, $query_data)
$query.wait_for_full_response do
	handle_socket($sock)
end

$query.servers.each do |s|
	puts s.inspect
end