xdiff_dbaction.py

#!/usr/bin/env python
#
# Copyright (C) 2018  Fernando Arnaboldi
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
from __future__ import print_function
import getopt
import itertools
import os.path
import sys
import time
import classes.settings
import classes.compat
from classes.dump import Dump
from classes.dbsqlite import DbSqlite


class Dbaction(object):
	"""Do stuff with the fuzzer's databases: copy databases, print tables, insert stuff and generate testcases"""
	def __init__(self, settings):
		self.settings = settings
		if 'max_permutation' not in self.settings:
			self.settings['max_permutation'] = 5
		if 'generate_type' not in self.settings:
			self.settings['generate_type'] = 2

	def print_table(self, fromdb, table, output_type):
		"""Print all the conents of a table"""
		if table is None:
			self.settings['logger'].error("You must select a table.")
		else:
			self.settings['output_file'] = None
			self.settings['db'] = DbSqlite(self.settings, fromdb)
			columns = self.settings['db'].get_columns(table)
			rows = self.settings['db'].get_rows(table)
			if columns:
				dump = Dump(self.settings)
				dump.general(output_type, table, columns, [rows])
			else:
				self.print_valid_tables(table)

	def insert_table(self, fromdb, table, separator, insert):
		"""Insert a row into a table"""
		if table is None:
			self.settings['logger'].error("You must select a table.")
		else:
			if not insert:
				self.settings['logger'].error("There are no values to be inserted")
			else:
				self.settings['db'] = DbSqlite(self.settings, fromdb)
				columns = self.settings['db'].get_columns(table)
				if columns:
					# If the user supplied one value less than the one required and the first column is called id, just ignore that column..
					if len(columns) == (len(insert.split(separator)) + 1) and columns[0] == 'id':
						del columns[0]
					if len(columns) != len(insert.split(separator)):
						print("The table '" + table + "' has " + str(len(columns)) + " columns: " + str(columns) + ". However, you want to insert " + str(len(insert.split(separator))) + " value/s: " + str(insert.split(separator)) + ". It doesn't work like that.")
					else:
						self.settings['db'].insert_row(table, columns, insert.split(separator))
				else:
					self.print_valid_tables(table)

	def print_valid_tables(self, table=None):
		"""Provide information on what are the valid tables"""
		if table:
			self.settings['logger'].error("Error: table '%s' not found" % table)
		else:
			if self.output_type:
				print("Valid table names:")
				print("- fuzz_testcase: contains the inputs to be sent to the software. You can define an input in 'function' and potential values in 'value' and generate the combinations on this table.")
				print("- function: contains what you want to fuzz. The special keyword [[test]] gets replaced by the values contained in the table 'value'. Ie, if you want to fuzz the 'print()'' function, you want to write in here 'print([[test]])'.")
				print("- value: contains the items that will replace the [[test]] values in the 'function' table")
				print("")
				print("Valid tables generated by XDiFF:")
				print("- fuzz_software: contains the software defined in software.ini")
				print("- fuzz_testcase_result: contains the result of executing the software defined in 'fuzz_software' with the input defined in 'fuzz_testcase'")
				print("- fuzz_constants: contains internal constant values used by the fuzzer")

	def permute(self, functions, values):
		"""Perform a permutation between the two lists received (functions & values)"""
		total = 0
		if not functions:
			self.settings['logger'].error("There are no functions to permute")
		elif not values:
			self.settings['logger'].error("There are no values to permute")
		else:
			# Prioritize the lower count injections
			for count in range(0, self.settings['max_permutation'] + 1):
				# Give a heads up of how many testcases will be generated
				subtotal = 0
				countfunctions = functions
				for function in countfunctions:
					if isinstance(function, tuple):
						if len(function) == 1:
							function = function[0]  # when it is generated by random testcases (classes/fuzzer.py)
						elif len(function) == 2:
							function = function[1]  # when it is read from the database
					if function is not None and count == function.count("[[test]]"):
						subtotal += 1
				self.settings['logger'].debug("Testcases generation: %s entry points, %s testcases to be generated." % (str(count), str(subtotal)))
				# Generate the testcases
				for function in functions:
					if len(function) == 1:
						function = function[0]  # when it is generated by random testcases (classes/fuzzer.py)
					elif len(function) == 2:
						function = function[1]  # when it is read from the database
					if function is not None and count == function.count("[[test]]"):
						testcases, total = self.permute_values(values, function, total)
						self.settings['db'].set_testcase(testcases)
		return total

	def permute_values(self, values, function, total):
		"""Perform a permutation between the values and the functions received based on the generate_type received"""
		testcases = []
		function_tuple = function

		# There are no values, only functions:
		if not values:
			testcases.append((classes.compat.unicode(function_tuple),))
		else:
			if self.settings['generate_type'] == 1:
				# Permute
				for valuetuple in itertools.product(values, repeat=function_tuple.count("[[test]]")):
					total += 1
					for value in valuetuple:
						# unicode values are tuples
						if isinstance(valuetuple, tuple):
							value = value[0]
						value = value.replace('[[id]]', str(total))
						function_tuple = function_tuple.replace("[[test]]", value, 1)
					testcases.append((classes.compat.unicode(function_tuple),))
					function_tuple = function  # reset to the original value
			elif self.settings['generate_type'] == 2:
				# Do not permute, just replace
				for value in values:
					if isinstance(value, tuple):
						value = value[0]
					total += 1
					value = value.replace('[[id]]', str(total))
					function_tuple = function_tuple.replace('[[test]]', value)
					testcases.append((classes.compat.unicode(function_tuple),))
					function_tuple = function  # reset to the original value
			elif self.settings['generate_type'] == 3:
				# Do not permute, replace but also include testcases with less parameters
				if (function.count("[[test]]")) > 1:
					for tests in range(1, function.count("[[test]]") + 1):
						for value in values:
							if isinstance(value, tuple):
								value = value[0]
							total += 1
							value = value.replace('[[id]]', str(total))
							function_tuple = function_tuple.replace('[[test]]', value)
							testcases.append((classes.compat.unicode(function_tuple),))
							function_tuple = function  # reset to the original value
						function_tuple = function = function.replace(',[[test]]', '', 1)
			else:
				print("Error: the permutation type does not exist")
				sys.exit()

		return testcases, total

	def generate(self, fromdb):
		"""Generate the testcases with a permutation of values and functions"""
		start_time = time.time()
		self.settings['db'] = DbSqlite(self.settings, fromdb)
		if self.settings['db'].db_connection:
			self.settings['db'].create_table()
			values = self.settings['db'].get_values()
			functions = self.settings['db'].get_functions()
			self.settings['logger'].info("Values: %s - Functions: %s" % (str(len(values)), str(len(functions))))
			total = self.permute(functions, values)

			self.settings['db'].commit()
			finish_time = time.time() - start_time
			self.settings['logger'].info("Testcases generated: %s" % str(total))
			self.settings['logger'].info("Time required: %s seconds" % str(round(finish_time, 2)))

	def migrate(self, fromdb, todb):
		"""Migrates tables from one database ('dbfrom') to another database ('dbto')"""
		start_time = time.time()
		self.settings['dbfrom'] = DbSqlite(self.settings, fromdb)
		self.settings['dbto'] = DbSqlite(self.settings, todb)

		if self.settings['dbfrom'].db_connection and self.settings['dbto'].db_connection:
			self.settings['dbto'].create_table()

			values = self.settings['dbfrom'].get_values()
			self.settings['dbto'].set_values(values)

			functions = self.settings['dbfrom'].get_functions()
			self.settings['dbto'].set_functions(functions)

			self.settings['dbto'].commit()

			finish_time = time.time() - start_time
			self.settings['logger'].info("Finished, time elapsed %s seconds" % str(finish_time)[:5])


def help(err=None):
	"""Print a help screen and exit"""
	if err:
		print("Error: %s\n" % str(err))
	print("Syntax: ")
	print(os.path.basename(__file__) + " -d db.sqlite -D fuzz.db             Migrate values and functions to another database")
	print("\t\t  -d fuzz.db -g 1 [-m 5]              Generate testcases permuting values and functions (set to maximum 5 input test cases)")
	print("\t\t  -d fuzz.db -g 2 [-m 5]              Generate testcases replacing values in functions (set to max..)")
	print("\t\t  -d fuzz.db -g 3 [-m 5]              Generate testcases replacing values in functions including testcases with less parameters (set to max..)")
	print("\t\t  -d fuzz.db -t table -p              Print a database table: fuzz_software, fuzz_testcase, value, function)")
	print("\t\t  -d fuzz.db -t table [-s,] -i \"foo\"  Insert foo into table (optional field separator -s uses a comma)")
	sys.exit()


def main():
	"""Perform multiple database actions"""
	try:
		opts, args = getopt.getopt(sys.argv[1:], "hd:D:g:i:m:ps:t:", ["help", "database=", "Database=", "generate=", "insert=", "maximum=", "print", "separator=", "table="])
	except getopt.GetoptError as err:
		help(err)

	settings = {}
	settings['output_type'] = 'txt'
	fromdb = None
	todb = None
	table = None
	action = None
	separator = ","

	for o, a in opts:
		if o in ("-h", "--help"):
			help()
		elif o in ("-d", "--database"):
			fromdb = a
			if os.path.isfile(fromdb):
				settings['db_file'] = fromdb
			else:
				help("The database selected '%s' is not a valid file." % a)
		elif o in ("-D", "--Database"):
			todb = a
			action = "migrate"
			break
		elif o in ("-g", "--generate"):
			action = "generate"
			try:
				settings['generate_type'] = int(a)
			except:
				help("The generate parameter should be a number")
		elif o in ("-i", "--insert"):
			action = "insert"
			insert = classes.compat.unicode(str(a), errors='ignore')
		elif o in ("-m", "--maximum"):
			try:
				settings['max_permutation'] = int(a)
			except ValueError:
				help("The max permutation parameter should be a number")
		elif o in ("-p", "--print"):
			action = "print"
		elif o in ("-s", "--separator"):
			separator = a
		elif o in ("-t", "--table"):
			table = a

	if not fromdb:
		help("The database was not specified.")

	settings = classes.settings.load_settings(settings)
	dbaction = Dbaction(settings)
	if action == "migrate":
		dbaction.migrate(fromdb, todb)
	elif action == "generate":
		if todb is not None:
			fromdb = todb
		dbaction.generate(fromdb)
	elif action == "print":
		dbaction.print_table(fromdb, table, settings['output_type'])
	elif action == "insert":
		dbaction.insert_table(fromdb, table, separator, insert)
	else:
		help("You must select an action: migrate, generate, print or insert.")


if __name__ == "__main__":
	main()