/nselib/socks.lua

---
-- A smallish SOCKS version 5 proxy protocol implementation
--
-- @author "Patrik Karlsson <patrik@cqure.net>"
--

local bin = require "bin"
local nmap = require "nmap"
local stdnse = require "stdnse"
local string = require "string"
_ENV = stdnse.module("socks", stdnse.seeall)

-- SOCKS Authentication methods
AuthMethod = {
	NONE		= 0,
	GSSAPI		= 1,
	USERPASS	= 2,
}

Request = {
	
	-- Class that handles the connection request to the server
	Connect = {
		
		-- Creates a new instance of the class
		-- @param auth_method table of requested authentication methods
		-- @return o instance on success, nil on failure
		new = function(self, auth_method)
			local o = { 
				version = 5,
				auth_method = ( "table" ~= type(auth_method) and { auth_method } or auth_method )
			}
			setmetatable(o, self)
			self.__index = self
			return o
		end,
		
		-- Converts the instance to string, so that it can be sent to the
		-- server.
		-- @return string containing the raw request
		__tostring = function(self)
			local methods = ""
			for _, m in ipairs(self.auth_method) do
				methods = methods .. string.char(m)
			end				
			return bin.pack("Cp", self.version, methods)
		end,	

	},
	
	-- Class that handles the authentication request to the server
	Authenticate = {
		
		-- Creates a new instance of the class
		-- @param auth_method number with the requested authentication method
		-- @param creds method specific table of credentials
		--        currently only user and pass authentication is supported
		--        this method requires two fields to be present
		--        <code>username</code> and <code>password</code>
		-- @return o instance on success, nil on failure
		new = function(self, auth_method, creds)
			local o = { 
				auth_method = auth_method,
				creds = creds
			}
			setmetatable(o, self)
			self.__index = self
			if ( auth_method == 2 ) then
				return o
			end
		end,
		
		-- Converts the instance to string, so that it can be sent to the
		-- server.
		-- @return string containing the raw request
		__tostring = function(self)
			-- we really don't support anything but 2, but let's pretend that
			-- we actually do
			if ( 2 == self.auth_method ) then
				local version = 1
				local username= self.creds.username or ""
				local password= self.creds.password or ""

				username = (username == "") and "\0" or username
				password = (password == "") and "\0" or password

				return bin.pack("Cpp", version, username, password)
			end
		end,
		
	}
	
}

Response = {
	
	-- Class that handles the connection response
	Connect = {
		
		-- Creates a new instance of the class
		-- @param data string containing the data as received over the socket
		-- @return o instance on success, nil on failure
		new = function(self, data)
			local o = { data = data }
			setmetatable(o, self)
			self.__index = self
			if ( o:parse() ) then
				return o
			end
		end,
		
		-- Parses the received data and populates member variables
		-- @return true on success, false on failure
		parse = function(self)
			if ( #self.data ~= 2 ) then
				return
			end
			local pos
			pos, self.version, self.method = bin.unpack("CC", self.data)	
			return true
		end
		
	},
	
	-- Class that handles the authentication response
	Authenticate = {
		
		Status = {
			SUCCESS = 0,
			-- could be anything but zero
			FAIL    = 1,
		},
		
		-- Creates a new instance of the class
		-- @param data string containing the data as received over the socket
		-- @return o instance on success, nil on failure
		new = function(self, data)
			local o = { data = data }
			setmetatable(o, self)
			self.__index = self
			if ( o:parse() ) then
				return o
			end
		end,
		
		-- Parses the received data and populates member variables
		-- @return true on success, false on failure
		parse = function(self)
			if ( #self.data ~= 2 ) then
				return
			end
			local pos
			pos, self.version, self.status = bin.unpack("CC", self.data)	
			return true
		end,
		
		-- checks if the authentication was successful or not
		-- @return true on success, false on failure
		isSuccess = function(self)
			return ( self.status == self.Status.SUCCESS )
		end,		
		
	}
	
}

-- A buffered socket implementation
Socket =
{	
	retries = 3,
	
	-- Creates a new socket instance
	-- @param host table containing the host table
	-- @param port table containing the port table
	-- @param options table containing options, currenlty supports:
	--        <code>timeout</code> - socket timeout in ms
	-- @return o new instance of Socket
	new = function(self, host, port, options)
		local o = { 
			host = host, 
			port = port,
			options = options or {}
		}
		setmetatable(o, self)
		self.__index = self
		o.Socket = nmap.new_socket()
		o.Buffer = nil
		return o
	end,
	
	-- Connects the socket to the server
	-- @return status true on success false on failure
	-- @return err string containing error message on failure
	connect = function( self )
		self.Socket:set_timeout(self.options.timeout or 10000)
		return self.Socket:connect( self.host, self.port )
	end,
	
	-- Closes an open connection.
	--
	-- @return Status (true or false).
	-- @return Error code (if status is false).
	close = function( self )
		return self.Socket:close()
	end,
	
	-- Opposed to the <code>socket:receive_bytes</code> function, that returns
	-- at least x bytes, this function returns the amount of bytes requested.
	--
	-- @param count of bytes to read
	-- @return true on success, false on failure
	-- @return data containing bytes read from the socket
	-- 		   err containing error message if status is false
	recv = function( self, count )
		local status, data
	
		self.Buffer = self.Buffer or ""
	
		if ( #self.Buffer < count ) then
			status, data = self.Socket:receive_bytes( count - #self.Buffer )
			if ( not(status) or #data < count - #self.Buffer ) then
				return false, data
			end
			self.Buffer = self.Buffer .. data
		end
			
		data = self.Buffer:sub( 1, count )
		self.Buffer = self.Buffer:sub( count + 1)
	
		return true, data	
	end,
	
	-- Sends data over the socket
	--
	-- @return Status (true or false).
	-- @return Error code (if status is false).
	send = function( self, data )
		return self.Socket:send( data )
	end,
}


-- The main script interface
Helper = {
	
	-- Create a new instance of the class
	-- @param host table containing the host table
	-- @param port table containing the port table
	-- @param options table containing library options, currenlty:
	--        <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 }
		setmetatable(o, self)
		self.__index = self
		return o
	end,
	
	-- Get the authentication method name by number
	-- @param method number containing the authentication method
	-- @return string containing the method name or Unknown
	authNameByNumber = function(self, method)
		local methods = {
			[0]  = "No authentication",
			[1]  = "GSSAPI",
			[2]  = "Username and password",
		}
		return methods[method] or ("Unknown method (%d)"):format(method)
	end,
	
	-- Connects to the SOCKS server
	-- @param auth_method table containing the auth. methods to request
	-- @return status true on success, false on failure
	-- @return response table containing the respons or err string on failure
	connect = function(self, auth_method)
		self.socket = Socket:new(self.host, self.port, self.options)		
		local status, err = self.socket:connect()
		if ( not(status) ) then
			return status, err
		end
		
		auth_method = auth_method or {AuthMethod.NONE, AuthMethod.GSSAPI, AuthMethod.USERPASS}
		status = self.socket:send( tostring(Request.Connect:new(auth_method)) )
		if ( not(status) ) then
			self.socket:close()
			return false, "Failed to send connection request to server"
		end
		
		local status, data = self.socket:recv(2)
		if ( not(status) ) then
			self.socket:close()
			return false, "Failed to receive connection response from server"
		end
		
		local response = Response.Connect:new(data)
		if ( not(response) ) then
			return false, "Failed to parse response from server"
		end

		if ( response.version ~= 5 ) then
			return false, ("Unsupported SOCKS version (%d)"):format(response.version)
		end
		if ( response.method == 0xFF ) then
			return false, "No acceptable authentication methods"
		end
		
		-- store the method so authenticate knows what to use
		self.auth_method = response.method
		return true, response
	end,
	
	-- Authenticates to the SOCKS server
	-- @param creds table containing authentication method specific fields
	--        currently only authentication method 2 (username and pass) is
	--        implemented. That method requires the following fields:
	--        <code>username</code> - containing the username
	--        <code>password</code> - containing the password
	-- @return status true on success, false on failure
	-- @return err string containing the error message
	authenticate = function(self, creds)
		if ( self.auth_method ~= 2 ) then
			return false, "Authentication method not supported"
		end
		local req = Request.Authenticate:new(self.auth_method, creds)
		if ( not(req) ) then
			return false, "Failed to create authentication request"
		end
		
		local status = self.socket:send(tostring(req))
		if ( not(status) ) then
			return false, "Failed to send authentication request"
		end
		
		if ( 2 == self.auth_method ) then
			local status, data = self.socket:recv(2)
			local auth = Response.Authenticate:new(data)
			
			if ( not(auth) ) then
				return false, "Failed to parse authentication response"
			end
			
			if ( auth:isSuccess() ) then
				return true, "Authentication was successfull"
			else
				return false, "Authentication failed"
			end

		end
		return false, "Unsupported authentication method"
	end,

	-- closes the connection to the server
	close = function(self)
		return self.socket:close()
	end,
	
}

return _ENV;