#!/usr/bin/env ruby
#
# This plugin provides integration with OpenVAS. Written by kost and
# averagesecurityguy.
#
# Distributed under MIT license: 
# http://www.opensource.org/licenses/mit-license.php
#

require 'socket' 
require 'timeout'
require 'openssl'
require 'rexml/document'
require 'rexml/text'
require 'base64'

# OpenVASOMP module
# 
# Usage: require 'openvas-omp'

module OpenVASOMP

#------------------------------
# Error Classes
#------------------------------
	class OMPError < :: RuntimeError
		attr_accessor :reason
		def initialize(reason = '')
			self.reason = reason
		end
		def to_s
			"OpenVAS OMP: #{self.reason}"
		end
	end

	class OMPConnectionError < OMPError
		def initialize
			self.reason = "Could not connect to server"
		end
	end

	class OMPResponseError < OMPError
		def initialize
			self.reason = "Error in OMP request/response"
		end
	end

	class OMPAuthError < OMPError
		def initialize
			self.reason = "Authentication failed"
		end
	end

	class XMLParsingError < OMPError
		def initialize
			self.reason = "XML parsing failed"
		end
	end


#------------------------------
# Connection Class
#------------------------------
	class OpenVASConnection
		attr_accessor :socket, :bufsize, :debug

		def initialize(host="127.0.0.1", port=9390, debug=false)
			@host = host
			@port = port
			@socket = nil
			@bufsize = 16384
			@debug = debug
		end

		def connect
			if @debug then puts "Connecting to server #{@host} on port #{@port}" end
			plain_socket = TCPSocket.open(@host, @port)
			ssl_context = OpenSSL::SSL::SSLContext.new()
			@socket = OpenSSL::SSL::SSLSocket.new(plain_socket, ssl_context)
			@socket.sync_close = true
			@socket.connect
		end
			
		def disconnect
			if @debug then
				puts "Closing connection to server #{@host} on port #{@port}" end
			if @socket then @socket.close end
		end

		def sendrecv(data)
			# Send the data
			if @debug then puts "Preparing to send data" end
			if not @socket then connect end
			if @debug then puts "SENDING: " + data end
			@socket.puts(data)

			# Receive the response
			resp = ''
			size = 0
			begin	
				begin
					timeout(@read_timeout) {
			    		a = @socket.sysread(@bufsize)
			    		size = a.length
			    		resp << a
					}
				rescue Timeout::Error
					size = 0
				rescue EOFError
					raise OMPResponseError
				end
			end while size >= @bufsize
			
			if @debug then puts "RECEIVED: " + resp end
			return resp
		end
	end
		

#------------------------------
# OpenVASOMP class
#------------------------------
	class OpenVASOMP
		# initialize object: try to connect to OpenVAS using URL, user and password
		attr_reader :targets, :tasks, :configs, :formats, :reports

		def initialize(user="openvas", pass="openvas", host="localhost", port=9392, debug=false)
			@debug = debug
			@token = ''
			@server = OpenVASConnection.new(host, port, debug)
			@server.connect
			login(user, pass)
			@configs = nil
			@tasks = nil
			@targets = nil
			@formats = nil
			@reports = nil
			config_get_all
			task_get_all
			target_get_all
			format_get_all
			report_get_all
		end

	#--------------------------
	# Low level commands. Only
	# used by OpenVASOMP class.
	#--------------------------
		# Nests a string inside an XML tag specified by root
		def xml_str(root, str)
			return "<#{root}>#{str}</#{root}>"
		end

		# Creates an XML root with child elements specified by a hash
		def xml_elems(root, elems)
			xml = REXML::Element.new(root)
			elems.each do |key, val|
				e = xml.add_element(key)
				e.text = val
			end
			return xml.to_s
		end

		# Creates and XML element with attributes specified by a hash
		def xml_attrs(elem, attribs)
			xml = REXML::Element.new(elem)
			attribs.each do |key, val|
				xml.attributes[key] = val
			end
			return xml.to_s
		end

		# Send authentication string and return an XML object (authentication token)
		def auth_request_xml(request)
			if @debug
				puts "Sending Request: #{request}"
			end
			resp = @server.sendrecv(request)
			begin
				docxml = REXML::Document.new(resp)
				status = docxml.root.attributes['status'].to_i
				status_text = docxml.root.attributes['status_text']
				if @debug
					puts "Status: #{status}"
					puts "Status Text: #{status_text}"
				end
			rescue
				raise XMLParsingError
			end

			return status, status_text
		end

		# Send string request wrapped with authentication XML and return 
		# an XML object
		def omp_request_xml(request)
			if @debug
				puts "Sending Request: #{request}"
			end
			resp = @server.sendrecv(@token + request)
			begin
				# Wrap the response in XML tags to use next_element properly.
				docxml = REXML::Document.new("<response>" + resp + "</response>")
				resp = docxml.root.elements['authenticate_response'].next_element
				status = resp.attributes['status'].to_i
				status_text = resp.attributes['status_text']	
				if @debug
					puts "Status: #{status}"
					puts "Status Text: #{status_text}"
				end
			rescue
				raise XMLParsingError
			end

			return status, status_text, resp
		end

	#--------------------------
	# Class API methods.
	#--------------------------
		# Sets debug level
		def debug(value)
			if value == 0
				@debug = false
				@server.debug = false
				return "Debug is deactivated."
			else
				@debug = true
				@server.debug = true
				return "Debug is activated."
			end
		end

		# get OMP version (you don't need to be authenticated)
		def get_version
			status, status_text, resp = omp_request_xml("<get_version/>")
			begin
				version = resp.elements['version'].text
				return version
			rescue
				raise XMLParsingError 
			end
		end

		# login to OpenVAS server. 
		# if successful returns authentication XML for further usage
		# if unsuccessful returns empty string
		def login(user, pass)
			creds = xml_elems("credentials", {"username"=> user, "password" => pass})
			req = xml_str("authenticate", creds)
			status, status_text = auth_request_xml(req)

			if status == 200
				@token = req
			else
				raise OMPAuthError	
			end
		end

		# Logout by disconnecting from the server and deleting the
		# authentication string. There are no sessions in OMP, must
		# send the credentials every time.
		def logout
			@server.disconnect()
			@token = ''
		end

#------------------------------
# Target Functions
#------------------------------

		# OMP - Get all targets for scanning and returns array of hashes
		# with following keys: id,name,comment,hosts,max_hosts,in_use
		#
		# Usage:
		# array_of_hashes = target_get_all()
		# 
		def target_get_all()
			begin
				status, status_text, resp = omp_request_xml("<get_targets/>")

				list = Array.new
				resp.elements.each('//get_targets_response/target') do |target|
					td = Hash.new
					td["id"] = target.attributes["id"]
					td["name"] = target.elements["name"].text
					td["comment"] = target.elements["comment"].text
					td["hosts"] = target.elements["hosts"].text
					td["max_hosts"] = target.elements["max_hosts"].text
					td["in_use"] = target.elements["in_use"].text
					list.push td
				end
				@targets = list
				return list
			rescue 
				raise OMPResponseError
			end
		end

		# OMP - Create target for scanning
		#
		# Usage:
		#
		# target_id = ov.target_create("name"=>"localhost",
		# 	"hosts"=>"127.0.0.1","comment"=>"yes")
		# 
		def target_create(name, hosts, comment)
			req = xml_elems("create_target", {"name"=>name, "hosts"=>hosts, "comment"=>comment})

			begin
				status, status_text, resp = omp_request_xml(req)
				target_get_all
				return "#{status_text}: #{resp.attributes['id']}"
			rescue 
				raise OMPResponseError
			end
		end

		# OMP - Delete target 
		#
		# Usage:
		#
		# ov.target_delete(target_id)
		# 
		def target_delete(id) 
			target = @targets[id.to_i]
			if not target
				raise OMPError.new("Invalid target id.")
			end
			req = xml_attrs("delete_target",{"target_id" => target["id"]})
			begin
				status, status_text, resp = omp_request_xml(req)
				target_get_all
				return status_text
			rescue 
				raise OMPResponseError
			end
		end

	#--------------------------
	# Task Functions
	#--------------------------
		# In short: Create a task.
		#
		# The client uses the create_task command to create a new task. 
		# 
		def task_create(name, comment, config_id, target_id)
			config = @configs[config_id.to_i]
			target = @targets[target_id.to_i]
			config = xml_attrs("config", {"id"=>config["id"]})
			target = xml_attrs("target", {"id"=>target["id"]})
			namestr = xml_str("name", name)
			commstr = xml_str("comment", comment)
			
			req = xml_str("create_task", namestr + commstr + config + target)

			begin
				status, status_text, resp = omp_request_xml(req)
				task_get_all
				return "#{status_text}: #{resp.attributes['id']}"
			rescue 
				raise OMPResponseError
			end
		end

		# In short: Delete a task.
		#
		# The client uses the delete_task command to delete an existing task,
		# including all reports associated with the task. 
		# 
		def task_delete(task_id) 
			task = @tasks[task_id.to_i]
			if not task
				raise OMPError.new("Invalid task id.")
			end
			req = xml_attrs("delete_task",{"task_id" => task["id"]})
			begin
				status, status_text, resp = omp_request_xml(req)
				task_get_all
				return status_text
			rescue 
				raise OMPResponseError
			end
		end

		# In short: Get all tasks.
		#
		# The client uses the get_tasks command to get task information. 
		# 
		def task_get_all()
			begin
				status, status_text, resp = omp_request_xml("<get_tasks/>")
				
				list = Array.new
				resp.elements.each('//get_tasks_response/task') do |task|
					td = Hash.new
					td["id"] = task.attributes["id"]
					td["name"] = task.elements["name"].text
					td["comment"] = task.elements["comment"].text
					td["status"] = task.elements["status"].text
					td["progress"] = task.elements["progress"].text
					list.push td	
				end
				@tasks = list
				return list
			rescue
				raise OMPResponseError
			end
		end

		# In short: Manually start an existing task.
		#
		# The client uses the start_task command to manually start an existing
		# task.
		#
		def task_start(task_id)
			task = @tasks[task_id.to_i]
			if not task
				raise OMPError.new("Invalid task id.")
			end
			req = xml_attrs("start_task",{"task_id" => task["id"]})
			begin
				status, status_text, resp = omp_request_xml(req)
				return status_text
			rescue 
				raise OMPResponseError
			end
		end

		# In short: Stop a running task.
		#
		# The client uses the stop_task command to manually stop a running
		# task.
		#
		def task_stop(task_id) 
			task = @tasks[task_id.to_i]
			if not task
				raise OMPError.new("Invalid task id.")
			end
			req = xml_attrs("stop_task",{"task_id" => task["id"]})
			begin
				status, status_text, resp = omp_request_xml(req)
				return status_text
			rescue 
				raise OMPResponseError
			end
		end 

		# In short: Pause a running task.
		#
		# The client uses the pause_task command to manually pause a running
		# task.
		#
		def task_pause(task_id) 
			task = @tasks[task_id.to_i]
			if not task
				raise OMPError.new("Invalid task id.")
			end
			req = xml_attrs("pause_task",{"task_id" => task["id"]})
			begin
				status, status_text, resp = omp_request_xml(req)
				return status_text
			rescue 
				raise OMPResponseError
			end
		end

		# In short: Resume task if stopped, else start task.
		#
		# The client uses the resume_or_start_task command to manually start
		# an existing task, ensuring that the task will resume from its
		# previous position if the task is in the Stopped state.
		# 
		def task_resume_or_start(task_id) 
			task = @tasks[task_id.to_i]
			if not task
				raise OMPError.new("Invalid task id.")
			end
			req = xml_attrs("resume_or_start_task",{"task_id" => task["id"]})
			begin
				status, status_text, resp = omp_request_xml(req)
				return status_text
			rescue 
				raise OMPResponseError
			end
		end

		# In short: Resume a puased task
		#
		# The client uses the resume_paused_task command to manually resume
		# a paused task.
		# 
		def task_resume_paused(task_id) 
			task = @tasks[task_id.to_i]
			if not task
				raise OMPError.new("Invalid task id.")
			end
			req = xml_attrs("resume_paused_task",{"task_id" => task["id"]})
			begin
				status, status_text, resp = omp_request_xml(req)
				return status_text
			rescue 
				raise OMPResponseError
			end
		end

	#--------------------------
	# Config Functions
	#--------------------------
		# OMP - get configs and returns hash as response
		# hash[config_id]=config_name
		#
		# Usage:
		#
		# array_of_hashes=ov.config_get_all()
		# 
		def config_get_all()
			begin
				status, status_text, resp = omp_request_xml("<get_configs/>")

				list = Array.new
				resp.elements.each('//get_configs_response/config') do |config|
					c = Hash.new
					c["id"] = config.attributes["id"]
					c["name"] = config.elements["name"].text
					list.push c
				end
				@configs = list
				return list
			rescue 
				raise OMPResponseError
			end
		end	


	#--------------------------
	# Format Functions
	#--------------------------
		# Get a list of report formats
		def format_get_all()
			begin
				status, status_text, resp = omp_request_xml("<get_report_formats/>")
				if @debug then print resp end
				
				list = Array.new
				resp.elements.each('//get_report_formats_response/report_format') do |report|
					td = Hash.new
					td["id"] = report.attributes["id"]
					td["name"] = report.elements["name"].text
					td["extension"] = report.elements["extension"].text
					td["summary"] = report.elements["summary"].text
					list.push td	
				end
				@formats = list
				return list
			rescue
				raise OMPResponseError
			end
		end


	#--------------------------
	# Report Functions
	#--------------------------
		# Get a list of reports
		def report_get_all()
			begin
				status, status_text, resp = omp_request_xml("<get_reports/>")
				
				list = Array.new
				resp.elements.each('//get_reports_response/report') do |report|
					td = Hash.new
					td["id"] = report.attributes["id"]
					td["task"] = report.elements["report/task/name"].text
					td["start_time"] = report.elements["report/scan_start"].text
					td["stop_time"] = report.elements["report/scan_end"].text
					list.push td	
				end
				@reports = list
				return list
			rescue
				raise OMPResponseError
			end
		end

		def report_delete(report_id) 
			report = @reports[report_id.to_i]
			if not report
				raise OMPError.new("Invalid report id.")
			end
			req = xml_attrs("delete_report",{"report_id" => report["id"]})
			begin
				status, status_text, resp = omp_request_xml(req)
				report_get_all
				return status_text
			rescue 
				raise OMPResponseError
			end
		end

		# Get a report by id. Must also specify the format_id
		def report_get_by_id(report_id, format_id)
			report = @reports[report_id.to_i]
			if not report
				raise OMPError.new("Invalid report id.")
			end
			format = @formats[format_id.to_i]
			if not format
				raise OMPError.new("Invalid format id.")
			end
			req = xml_attrs("get_reports", {"report_id"=>report["id"], "format_id"=>format["id"]})

			begin
				status, status_text, resp = omp_request_xml(req)
				if status == "404"
					raise OMPError.new(status_text)
				end
				content_type = resp.elements["report"].attributes["content_type"]
				report = resp.elements["report"].text

				if report == nil
					raise OMPError.new("The report is empty.")
				end

				# XML reports are in XML format, everything else is base64 encoded.
				if content_type == "text/xml"
					return resp.elements["report"].text
				else
					return Base64.decode64(resp.elements["report"].text)
				end
			rescue
				raise OMPResponseError
			end
		end

    end
end
