Code Refactoring #2

src/forbidden/main.py

@@ -102,15 +102,7 @@ def get_encoded_domains(dnp, port):
 
 # ----------------------------------------
 
-path_const = "/"
 
-def replace_multiple_slashes(path):
-	return re.sub(r"\/{2,}", path_const, path)
-
-def prepend_slash(path):
-	if not path.startswith(path_const):
-		path = path_const + path
-	return path
 
 def append_paths(bases, paths):
 	if not isinstance(bases, list):
@@ -299,7 +291,7 @@ def __eq__(self, other):
 	def lower(self):
 		if self.__lower is None:
 			lower = str.lower(self)
-			if str.__eq__(lower, self): 
+			if str.__eq__(lower, self):
 				self.__lower = self
 			else:
 				self.__lower = uniquestr(lower)

src/forbidden/utils/cookie.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+
+from . import grep
+
+def get_key_value(cookie: str):
+	"""
+	Get a key-value pair from an HTTP cookie.\n
+	Returns an empty key-value pair on failure.
+	"""
+	key = ""; value = ""
+	if grep.search(r"^[^\=\;]+\=[^\=\;]+$|^[^\=\;]+\=$", cookie):
+		key, value = cookie.split("=", 1)
+	return key.strip(), value.strip()
+
+def format_key_value(key: str, value: str):
+	"""
+	Returns a key-value pair as a string.
+	"""
+	return f"{key}={value}"

src/forbidden/utils/file.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+
+from . import array
+
+import os
+
+__ENCODING = "ISO-8859-1"
+
+def validate(file: str):
+	"""
+	Validate a file.\n
+	Success flag is 'True' if the 'file' exists and is a regular file, has a read permission and is not empty.
+	"""
+	success = False
+	message = ""
+	if not os.path.isfile(file):
+		message = f"\"{file}\" does not exist"
+	elif not os.access(file, os.R_OK):
+		message = f"\"{file}\" does not have a read permission"
+	elif not os.stat(file).st_size > 0:
+		message = f"\"{file}\" is empty"
+	else:
+		success = True
+	return success, message
+
+def read_array(file: str) -> list[str]:
+	"""
+	Read a file line by line, and append the lines to a list.\n
+	Whitespace will be stripped from each line, and empty lines will be removed.\n
+	Returns a unique list.
+	"""
+	tmp = []
+	with open(file, "r", encoding = __ENCODING) as stream:
+		for line in stream:
+			line = line.strip()
+			if line:
+				tmp.append(line)
+	return array.unique(tmp)

src/forbidden/utils/general.py

@@ -1,5 +1,54 @@
 #!/usr/bin/env python3
 
+import enum
+
+class Test(enum.Enum):
+	"""
+	Enum containing supported tests.
+	"""
+	BASE             = "base"
+	METHODS          = "methods"
+	METHOD_OVERRIDES = "method-overrides"
+	SCHEME_OVERRIDES = "scheme-overrides"
+	PORT_OVERRIDES   = "port-overrides"
+	HEADERS          = "headers"
+	VALUES           = "values"
+	PATHS            = "paths"
+	PATHS_RAM        = "paths-ram"
+	ENCODINGS        = "encodings"
+	AUTHS            = "auths"
+	REDIRECTS        = "redirects"
+	PARSERS          = "parsers"
+
+	@classmethod
+	def all(cls):
+		"""
+		Get all supported tests.
+		"""
+		return [
+			cls.BASE,
+			cls.METHODS,
+			cls.METHOD_OVERRIDES,
+			cls.SCHEME_OVERRIDES,
+			cls.PORT_OVERRIDES,
+			cls.HEADERS,
+			cls.VALUES,
+			cls.PATHS,
+			cls.PATHS_RAM,
+			cls.ENCODINGS,
+			cls.AUTHS,
+			cls.REDIRECTS,
+			cls.PARSERS
+		]
+
+# ----------------------------------------
+
+PATHS = ["/robots.txt", "/index.html", "/sitemap.xml", "/README.txt"]
+
+EVIL_URL = "https://github.com"
+
+# ----------------------------------------
+
 def print_error(message: str):
 	"""
 	Print an error message.

src/forbidden/utils/grep.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python3
+
+import regex as re
+
+__FLAGS = re.MULTILINE | re.IGNORECASE
+
+def validate(query: str):
+	"""
+	Validate a regular expression.
+	"""
+	success = False
+	message = ""
+	try:
+		re.compile(query)
+		success = True
+	except re.error:
+		message = f"Invalid RegEx: {query}"
+	return success, message
+
+def search(string: str, query: str):
+	"""
+	Check if there are any matches in a string using the specified RegEx pattern.
+	"""
+	return bool(re.search(query, string, flags = __FLAGS))

src/forbidden/utils/header.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+
+from . import grep
+
+def get_key_value(header: str):
+	"""
+	Get a key-value pair from an HTTP request header.\n
+	Returns an empty key-value pair on failure.
+	"""
+	key = ""; value = ""
+	if grep.search(r"^[^\:]+\:.+$", header):
+		key, value = header.split(":", 1)
+	elif grep.search(r"^[^\;]+\;$", header):
+		key, value = header.split(";", 1)
+	return key.strip(), value.strip()
+
+def format_key_value(key: str, value: str):
+	"""
+	Returns a key-value pair as a string.
+	"""
+	return f"{key}: {value}" if value else f"{key};"

src/forbidden/utils/path.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+
+import regex as re
+
+__SEP = "/"
+
+def replace_multiple_slashes(path: str):
+	"""
+	Replace multiple consecutive forward slashes with a single forward slash.
+	For example, replace '//' with '/', etc.
+	"""
+	return re.sub(r"\/{2,}", __SEP, path)
+
+def prepend_slash(path: str):
+	"""
+	Append a single forward slash if one does not already exist.
+	"""
+	if not path.startswith(__SEP):
+		path = __SEP + path
+	return path

src/forbidden/utils/validate.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-from . import array, config, general, test, url
+from . import array, config, cookie, file, general, header, path, url
 
 import argparse, sys
 
@@ -182,18 +182,63 @@ def __validate_url(self):
 			self.__error(f"Inaccessible URL: {message}")
 
 	def __validate_tests(self):
-		types = array.remove_empty_strings(self.__args.tests.split(","))
+		tests = array.remove_empty_strings(self.__args.tests.split(","))
 		self.__args.tests = []
-		if not types:
+		if not tests:
 			self.__error("No tests were specified")
 		else:
-			supported = [type.value for type in test.Type.types()]
-			for type in types:
-				if type == "all":
-					self.__args.tests.extend(test.Type.all())
-				elif type not in supported:
+			supported = [test.value for test in general.Test.all()]
+			for test in tests:
+				if test == "all":
+					self.__args.tests.extend(general.Test.all())
+				elif test not in supported:
 					self.__error("Supported tests are 'base', 'methods', '(method|scheme|port)-overrides', 'headers', 'values', 'paths(-ram)', 'encodings', 'auths', 'redirects', 'parsers', or 'all'")
 					break
 				else:
-					self.__args.tests.append(test.Type(type))
-			self.__args.tests = array.unique(self.__args.tests)
+					self.__args.tests.append(general.Test(test))
+		self.__args.tests = array.unique(self.__args.tests)
+
+	def __validate_values(self):
+		tmp = []
+		if self.__args.values:
+			success, message = file.validate(self.__args.values)
+			if not success:
+				self.__error(message)
+			else:
+				tmp = file.read_array(self.__args.values)
+				if not tmp:
+					self.__error(f"No values were found in \"{self.__args.values}\"")
+		self.__args.values = tmp
+
+	def __validate_path(self):
+		self.__args.path = [path.prepend_slash(path.replace_multiple_slashes(self.__args.path))] if self.__args.path else general.PATHS
+
+	def __validate_evil(self):
+		if self.__args.evil:
+			success, message = url.validate(self.__args.evil)
+			if not success:
+				self.__error(f"Evil URL: {message}")
+		else:
+			self.__args.evil = general.EVIL_URL
+
+	def __validate_header(self):
+		tmp = []
+		if self.__args.header:
+			for entry in self.__args.header:
+				key, value = header.get_key_value(entry[0])
+				if not key:
+					self.__error(f"Invalid HTTP request header: {entry[0]}")
+					continue
+				tmp.append(header.format_key_value(key, value))
+		self.__args.header = tmp
+
+	def __validate_cookie(self):
+		tmp = []
+		if self.__args.cookie:
+			for entry in self.__args.cookie:
+				key, value = cookie.get_key_value(entry[0])
+				if not key:
+					self.__error(f"Invalid HTTP cookie: {entry[0]}")
+					continue
+				tmp.append(cookie.format_key_value(key, value))
+		self.__args.cookie = tmp

src/stresser/main.py

@@ -175,7 +175,7 @@ def __eq__(self, other):
 	def lower(self):
 		if self.__lower is None:
 			lower = str.lower(self)
-			if str.__eq__(lower, self): 
+			if str.__eq__(lower, self):
 				self.__lower = self
 			else:
 				self.__lower = uniquestr(lower)