---
-- A minimalistic Asterisk IAX2 (Inter-Asterisk eXchange v2) VoIP protocol implementation.
-- The library implements the minimum needed to perform brute force password guessing.
--
-- @author "Patrik Karlsson <patrik@cqure.net>"
--

local bin = require "bin"
local bit = require "bit"
local math = require "math"
local nmap = require "nmap"
local os = require "os"
local stdnse = require "stdnse"
local openssl = stdnse.silent_require "openssl"
local table = require "table"
_ENV = stdnse.module("iax2", stdnse.seeall)


IAX2 = {
	
	FrameType = {
		IAX = 6,		
	},
	
	SubClass = {
		ACK 		= 0x04,
		REGACK		= 0x0f,
		REGREJ		= 0x10,
		REGREL 		= 0x11,
		CALLTOKEN 	= 0x28,
	},
	
	InfoElement = {
		USERNAME 	= 0x06,	
		CHALLENGE	= 0x0f,
		MD5_RESULT	= 0x10,
		CALLTOKEN   = 0x36,
	},
	
	PacketType = {
		FULL = 1,
	},
	
	-- The IAX2 Header
	Header = {
		
		-- Creates a new Header instance
		-- @param src_call number containing the source call
		-- @param dst_call number containing the dest call
		-- @param timestamp number containing a timestamp
		-- @param oseqno number containing the seqno of outgoing packets
		-- @param iseqno number containing the seqno of incoming packets
		-- @param frametype number containing the frame type
		-- @param subclass number containing the subclass type
		new = function(self, src_call, dst_call, timestamp, oseqno, iseqno, frametype, subclass)
			local o = {
				type 		= IAX2.PacketType.FULL,
				retrans		= false,
				src_call 	= src_call,
				dst_call	= dst_call,
				timestamp	= timestamp,
				oseqno		= oseqno,
				iseqno		= iseqno,
				frametype	= frametype,
				subclass	= subclass,
			}
			setmetatable(o, self)
			self.__index = self
			return o
		end,
		
		-- Parses data, a byte string, and creates a new Header instance
		-- @return header instance of Header	
		parse = function(data)
			local header = IAX2.Header:new()
			local pos, frame_type = bin.unpack("C", data)
			if ( bit.band(frame_type, 0x80) == 0 ) then
				print("frame_type", stdnse.tohex(frame_type))
				stdnse.print_debug(2, "Frametype not supported")
				return
			end
			header.type = IAX2.PacketType.FULL
			pos, header.src_call = bin.unpack(">S", data)
			header.src_call = bit.band(header.src_call, 0x7FFF)

			local retrans
			pos, retrans = bin.unpack("C", data, pos)
			if ( bit.band(retrans, 0x80) == 8 ) then
				header.retrans = true
			end
			pos, header.dst_call = bin.unpack(">S", data, pos - 1)
			header.dst_call = bit.band(header.dst_call, 0x7FFF)
			
			pos, header.timestamp, header.oseqno, 
				header.iseqno, header.frametype, header.subclass = bin.unpack(">ICCCC", data, pos)
			
			return header
		end,
		
		-- Converts the instance to a string
		-- @return str containing the instance
		__tostring = function(self)
			assert(self.src_call < 32767, "Source call exceeds 32767")
			assert(self.dst_call < 32767, "Dest call exceeds 32767")
			local src_call = self.src_call
			local dst_call = self.dst_call
			if ( self.type == IAX2.PacketType.FULL ) then
				src_call = src_call + 32768
			end
			if ( self.retrans ) then
				dst_call = dst_call + 32768
			end
			return bin.pack(">SSICCCC", src_call, dst_call, self.timestamp,
				self.oseqno, self.iseqno, self.frametype, self.subclass)
		end,
	},
	
	-- The IAX2 Request class
	Request = {
		
		-- Creates a new instance
		-- @param header instance of Header
		new = function(self, header)
			local o = {
				header = header,
				ies = {}
			}
			setmetatable(o, self)
			self.__index = self
			return o			
		end,
		
		-- Sets an Info Element or adds one, in case it's missing
		-- @param key the key value of the IE to add
		-- @param value string containing the value to set or add
		setIE = function(self, key, value)
			for _, ie in ipairs(self.ies or {}) do
				if ( key == ie.type ) then
					ie.value = value
				end
			end
			table.insert(self.ies, { type = key, value = value } )
		end,
		
		-- Gets an information element
		-- @param key number containing the element number to retrieve
		-- @return ie table containing the info element if it exists
		getIE = function(self, key)
			for _, ie in ipairs(self.ies or {}) do
				if ( key == ie.type ) then
					return ie
				end
			end
		end,

		-- Converts the instance to a string
		-- @return str containing the instance
		__tostring = function(self)
			local data = ""
			for _, ie in ipairs(self.ies) do
				data = data .. bin.pack("Cp", ie.type, ie.value )
			end
			
			return tostring(self.header) .. data
		end,
		
	},
	
	
	-- The IAX2 Response
	Response = {
		
		-- Creates a new instance
		new = function(self)
			local o = { ies = {} }
			setmetatable(o, self)
			self.__index = self
			return o			
		end,
		
		-- Sets an Info Element or adds one, in case it's missing
		-- @param key the key value of the IE to add
		-- @param value string containing the value to set or add
		setIE = function(self, key, value)
			for _, ie in ipairs(self.ies or {}) do
				if ( key == ie.type ) then
					ie.value = value
				end
			end
			table.insert(self.ies, { type = key, value = value } )
		end,
		
		-- Gets an information element
		-- @param key number containing the element number to retrieve
		-- @return ie table containing the info element if it exists
		getIE = function(self, key)
			for _, ie in ipairs(self.ies or {}) do
				if ( key == ie.type ) then
					return ie
				end
			end
		end,
		
		-- Parses data, a byte string, and creates a response
		-- @return resp instance of response	
		parse = function(data)
			local resp = IAX2.Response:new()
			if ( not(resp) ) then return end

			resp.header = IAX2.Header.parse(data)
			if ( not(resp.header) ) then return	end
			
			local pos = 13
			resp.ies = {}
			repeat
				local ie = {}
				pos, ie.type, ie.value = bin.unpack(">Cp", data, pos)
				table.insert(resp.ies, ie)
			until( pos > #data )
			return resp
		end,
		
	}
		
}


Helper = {
	
	-- Creates a new Helper instance
	-- @param host table as received by the action method
	-- @param port table as received by the action method
	-- @param options table containing helper options, currently
	--        <code>timeout</code> socket timeout in ms
	-- @return o instance of Helper
	new = function(self, host, port, options)
		local o = { host = host, port = port, options = options or {} }
		setmetatable(o, self)
		self.__index = self
		return o			
	end,
	
	-- Connects the UDP socket to the server
	-- @return status true on success, false on failure
	-- @return err message containing error if status is false
	connect = function(self)
		self.socket = nmap.new_socket()
		self.socket:set_timeout(self.options.timeout or 5000)
		return self.socket:connect(self.host, self.port)
	end,
	
	-- Sends a request to the server and receives the response
	-- @param req instance containing the request to send to the server
	-- @return status true on success, false on failure
	-- @return resp instance of response on success,
	--         err containing the error message on failure
	exch = function(self, req)
		local status, err = self.socket:send(tostring(req))
		if ( not(status) ) then
			return false, "Failed to send request to server"
		end
		local status, data = self.socket:receive()
		if ( not(status) ) then
			return false, "Failed to receive response from server"
		end
		
		local resp = IAX2.Response.parse(data)
		return true, resp
	end,
	
	-- Request a session release
	-- @param username string containing the extention (username)
	-- @param password string containing the password
	regRelease = function(self, username, password)

		local src_call = math.random(32767)
		local header = IAX2.Header:new(src_call, 0, os.time(), 0, 0, IAX2.FrameType.IAX, IAX2.SubClass.REGREL)
		local regrel = IAX2.Request:new(header)

		regrel:setIE(IAX2.InfoElement.USERNAME, username)
		regrel:setIE(IAX2.InfoElement.CALLTOKEN, "")

		local status, resp = self:exch(regrel)
		if ( not(status) ) then
			return false, resp
		end

		if ( not(resp) or IAX2.SubClass.CALLTOKEN ~= resp.header.subclass ) then
			return false, "Unexpected response"
		end

		local token = resp:getIE(IAX2.InfoElement.CALLTOKEN)
		if ( not(token) ) then
			return false, "Failed to get token"
		end

		regrel:setIE(IAX2.InfoElement.CALLTOKEN, token.value)
		status, resp = self:exch(regrel)
		if ( not(status) ) then
			return false, resp
		end
		
		local challenge = resp:getIE(IAX2.InfoElement.CHALLENGE)
		if ( not(challenge) ) then
			return false, "Failed to retrieve challenge from server"
		end
		
		regrel.header.iseqno = 1
		regrel.header.oseqno = 1
		regrel.header.dst_call = resp.header.src_call
		regrel.ies = {}
		
		local hash = stdnse.tohex(openssl.md5(challenge.value .. password))
		regrel:setIE(IAX2.InfoElement.USERNAME, username)
		regrel:setIE(IAX2.InfoElement.MD5_RESULT, hash)
		
		status, resp = self:exch(regrel)
		if ( not(status) ) then
			return false, resp
		end
		
		if ( IAX2.SubClass.ACK == resp.header.subclass ) then
			local data
			status, data = self.socket:receive()		
			resp = IAX2.Response.parse(data)
		end
		
		if ( status and IAX2.SubClass.REGACK == resp.header.subclass ) then
			return true
		end
		return false, "Release failed"
	end,
		
	-- Close the connection with the server
	-- @return true on success, false on failure
	close = function(self)
		return self.socket:close()
	end,
	
	
}

return _ENV;
