New function and standarized methods

This commit is contained in:
Enrique Martin 2023-08-09 13:46:09 +02:00
parent 40421fadbb
commit 6d62c38ee1
10 changed files with 268 additions and 91 deletions

View File

@ -1,4 +1,6 @@
from .general import * from .general import *
from .output import *
from .encryption import *
from .threads import * from .threads import *
from .agents import * from .agents import *
from .modules import * from .modules import *

View File

@ -1,7 +1,5 @@
import sys import sys
import os import os
from .general import debug_dict,now,set_dict_key_value,generate_md5
from .modules import init_module,init_log_module,print_module,print_log_module
#### ####
# Define global variables dict, used in functions as default values. # Define global variables dict, used in functions as default values.
@ -43,6 +41,8 @@ def set_global_variable(
variable_name (str): Name of the variable to set. variable_name (str): Name of the variable to set.
value (any): Value to assign to the variable. value (any): Value to assign to the variable.
""" """
from .general import set_dict_key_value
set_dict_key_value(GLOBAL_VARIABLES, variable_name, value) set_dict_key_value(GLOBAL_VARIABLES, variable_name, value)
#### ####
@ -69,46 +69,47 @@ class Agent:
self.log_modules_def = log_modules_def self.log_modules_def = log_modules_def
self.added_modules = [] self.added_modules = []
'''
TODO: Add commnets
'''
def update_config( def update_config(
self, self,
config: dict = {} config: dict = {}
): ):
'''
TODO: Add commnets
'''
for key, value in config.items(): for key, value in config.items():
if key in self.config: if key in self.config:
self.config[key] = value self.config[key] = value
'''
TODO: Add commnets
'''
def get_config( def get_config(
self self
) -> dict: ) -> dict:
'''
TODO: Add commnets
'''
return self.config return self.config
'''
TODO: Add commnets
'''
def add_module( def add_module(
self, self,
module: dict = {} module: dict = {}
): ):
'''
TODO: Add commnets
'''
from .general import generate_md5
from .modules import init_module
if "name" in module and type(module["name"]) == str and len(module["name"].strip()) > 0: if "name" in module and type(module["name"]) == str and len(module["name"].strip()) > 0:
self.modules_def.append(init_module(module)) self.modules_def.append(init_module(module))
self.added_modules.append(generate_md5(module["name"])) self.added_modules.append(generate_md5(module["name"]))
'''
TODO: Add commnets
'''
def del_module( def del_module(
self, self,
module_name: str = "" module_name: str = ""
): ):
'''
TODO: Add commnets
'''
from .general import generate_md5
if len(module_name.strip()) > 0: if len(module_name.strip()) > 0:
try: try:
@ -120,15 +121,14 @@ class Agent:
self.added_modules.pop(module_id) self.added_modules.pop(module_id)
self.modules_def.pop(module_id) self.modules_def.pop(module_id)
'''
TODO: Add commnets
'''
def update_module( def update_module(
self, self,
module_name: str = "", module_name: str = "",
module: dict = {} module: dict = {}
): ):
'''
TODO: Add commnets
'''
module_def = self.get_module(module_name) module_def = self.get_module(module_name)
if module_def: if module_def:
@ -140,13 +140,14 @@ class Agent:
self.del_module(module_name) self.del_module(module_name)
self.add_module(module_def) self.add_module(module_def)
'''
TODO: Add commnets
'''
def get_module( def get_module(
self, self,
module_name: str = "" module_name: str = ""
) -> dict: ) -> dict:
'''
TODO: Add commnets
'''
from .general import generate_md5
if len(module_name.strip()) > 0: if len(module_name.strip()) > 0:
try: try:
@ -159,43 +160,41 @@ class Agent:
else: else:
return {} return {}
'''
TODO: Add commnets
'''
def get_modules_def( def get_modules_def(
self self
) -> dict: ) -> dict:
'''
TODO: Add commnets
'''
return self.modules_def return self.modules_def
'''
TODO: Add commnets
'''
def add_log_module( def add_log_module(
self, self,
log_module: dict = {} log_module: dict = {}
): ):
'''
TODO: Add commnets
'''
from .modules import init_log_module
if "source" in module and type(module["source"]) == str and len(module["source"].strip()) > 0: if "source" in module and type(module["source"]) == str and len(module["source"].strip()) > 0:
self.log_modules_def.append(init_log_module(log_module)) self.log_modules_def.append(init_log_module(log_module))
'''
TODO: Add commnets
'''
def get_log_modules_def( def get_log_modules_def(
self self
) -> dict: ) -> dict:
'''
TODO: Add commnets
'''
return self.log_modules_def return self.log_modules_def
'''
TODO: Add commnets
'''
def print_xml( def print_xml(
self, self,
print_flag: bool = False print_flag: bool = False
) -> str: ) -> str:
'''
TODO: Add commnets
'''
return print_agent(self.get_config(), self.get_modules_def(), self.get_log_modules_def(), print_flag) return print_agent(self.get_config(), self.get_modules_def(), self.get_log_modules_def(), print_flag)
#### ####
@ -210,6 +209,8 @@ def init_agent(
Returns: Returns:
dict: Dictionary representing the agent template with default values. dict: Dictionary representing the agent template with default values.
""" """
from .general import now
agent = { agent = {
"agent_name" : "", "agent_name" : "",
"agent_alias" : "", "agent_alias" : "",
@ -245,6 +246,9 @@ def print_agent(
- Use print_flag to show modules' XML in STDOUT. - Use print_flag to show modules' XML in STDOUT.
- Returns xml (str). - Returns xml (str).
""" """
from .output import print_stdout
from .modules import print_module,print_log_module
xml = "" xml = ""
data_file = None data_file = None
@ -267,6 +271,6 @@ def print_agent(
xml += "</agent_data>" xml += "</agent_data>"
if print_flag: if print_flag:
print(xml) print_stdout(xml)
return xml return xml

View File

@ -1,6 +1,5 @@
import sys import sys
import json import json
from .general import debug_dict
#### ####
# Define some global variables # Define some global variables
@ -137,6 +136,7 @@ def print_output():
to create the JSON output. It then prints the JSON string and exits the script with to create the JSON output. It then prints the JSON string and exits the script with
the 'ERROR_LEVEL' as the exit code. the 'ERROR_LEVEL' as the exit code.
""" """
from .output import print_stdout
global ERROR_LEVEL global ERROR_LEVEL
global SUMMARY global SUMMARY
@ -155,5 +155,5 @@ def print_output():
json_string = json.dumps(OUTPUT) json_string = json.dumps(OUTPUT)
print(json_string) print_stdout(json_string)
sys.exit(ERROR_LEVEL) sys.exit(ERROR_LEVEL)

View File

@ -0,0 +1,78 @@
try:
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
except ImportError as e:
import sys
from .output import print_stderr
print_stderr("ModuleNotFoundError: No module named 'pycryptodome'")
sys.exit(1)
import hashlib
import base64
import hmac
from binascii import unhexlify
####
# Define encription internal global variables.
#########################################################################################
_PASSWORD = "default_salt"
####
# Internal use only: Get AES cipher
#########################################################################################
def _get_cipher(
password: str = _PASSWORD
) -> AES:
'''
Internal use only: Get AES cipher
'''
key = b''
msg = password.encode('utf-8')
hash_obj = hmac.new(key, msg, hashlib.sha256)
hash_result = hash_obj.digest()
hash_base64 = base64.b64encode(hash_result)[:16].decode()
iv = b'0000000000000000'
return AES.new(hash_base64.encode(), AES.MODE_CBC, iv)
####
# Return encrypted string
#########################################################################################
def encrypt(
str_to_encrypt: str = "",
password: str = _PASSWORD
) -> str:
'''
Return encrypted string
'''
cipher = _get_cipher(password)
try:
msg_padded = pad(str_to_encrypt.encode(), AES.block_size, style='pkcs7')
cipher_text = cipher.encrypt(msg_padded)
b64str = base64.b64encode(cipher_text).decode()
except:
b64str = ''
return b64str
####
# Return decrypted string
#########################################################################################
def decrypt(
str_to_decrypt: str = "",
password: str = _PASSWORD
) -> str:
'''
Return decrypted string
'''
cipher = _get_cipher(password)
try:
decrypted_str = unpad(cipher.decrypt(base64.b64decode(str_to_decrypt)), AES.block_size, style='pkcs7').decode().strip()
except:
decrypted_str = ''
return decrypted_str

View File

@ -1,5 +1,4 @@
import sys import sys
import json
from datetime import datetime from datetime import datetime
import hashlib import hashlib
@ -323,24 +322,6 @@ def safe_output(
return input_string return input_string
####
# Prints dictionary in formatted json string.
#########################################################################################
def debug_dict(
jsontxt = ""
):
"""
Prints any list, dict, string, float or integer as a json
"""
try:
debug_json = json.dumps(jsontxt, indent=4)
print (debug_json)
except json.JSONDecodeError as e:
print(f"debug_dict: Failed to dump. Error: {e}")
except Exception as e:
print(f"debug_dict: Unexpected error: {e}")
#### ####
# Assign to a key in a dict a given value. # Assign to a key in a dict a given value.
######################################################################################### #########################################################################################
@ -386,13 +367,15 @@ def generate_md5(
######################################################################################### #########################################################################################
def now( def now(
print_flag: int = 0, print_flag: bool = False,
utimestamp: int = 0 utimestamp: bool = False
) -> str: ) -> str:
""" """
Returns time in yyyy/mm/dd HH:MM:SS format by default. Use 1 as an argument Returns time in yyyy/mm/dd HH:MM:SS format by default. Use 1 as an argument
to get epoch time (utimestamp) to get epoch time (utimestamp)
""" """
from .output import print_stdout
today = datetime.today() today = datetime.today()
if utimestamp: if utimestamp:
@ -401,7 +384,7 @@ def now(
time = today.strftime('%Y/%m/%d %H:%M:%S') time = today.strftime('%Y/%m/%d %H:%M:%S')
if print_flag: if print_flag:
print(time) print_stdout(time)
return time return time
@ -443,6 +426,8 @@ def parse_configuration(
Returns: Returns:
- dict: containing all keys and values from file. - dict: containing all keys and values from file.
""" """
from .output import print_stderr
config = {} config = {}
try: try:
@ -456,7 +441,7 @@ def parse_configuration(
config[option.strip()] = value.strip() config[option.strip()] = value.strip()
except Exception as e: except Exception as e:
print (f"{type(e).__name__}: {e}") print_stderr(f"{type(e).__name__}: {e}")
for option, value in default_values.items(): for option, value in default_values.items():
if option.strip() not in config: if option.strip() not in config:
@ -472,7 +457,7 @@ def parse_csv_file(
file: str = "", file: str = "",
separator: str = ';', separator: str = ';',
count_parameters: int = 0, count_parameters: int = 0,
debug: bool = False print_errors: bool = False
) -> list: ) -> list:
""" """
Parse csv configuration. Reads configuration file and stores its data in a list. Parse csv configuration. Reads configuration file and stores its data in a list.
@ -481,11 +466,13 @@ def parse_csv_file(
- file (str): configuration csv file path. \n - file (str): configuration csv file path. \n
- separator (str, optional): Separator for option and value. Defaults to ";". - separator (str, optional): Separator for option and value. Defaults to ";".
- coun_parameters (int): min number of parameters each line shold have. Default None - coun_parameters (int): min number of parameters each line shold have. Default None
- debug (bool): print errors on lines - print_errors (bool): print errors on lines
Returns: Returns:
- List: containing a list for of values for each csv line. - List: containing a list for of values for each csv line.
""" """
from .output import print_stderr
csv_arr = [] csv_arr = []
try: try:
@ -498,11 +485,11 @@ def parse_csv_file(
value = line.strip().split(separator) value = line.strip().split(separator)
if len(value) >= count_parameters: if len(value) >= count_parameters:
csv_arr.append(value) csv_arr.append(value)
elif debug==True: elif print_errors==True:
print(f'Csv line: {line} does not match minimun parameter defined: {count_parameters}',file=sys.stderr) print_stderr(f'Csv line: {line} does not match minimun parameter defined: {count_parameters}')
except Exception as e: except Exception as e:
print (f"{type(e).__name__}: {e}") print_stderr(f"{type(e).__name__}: {e}")
return csv_arr return csv_arr

View File

@ -2,7 +2,6 @@ from requests_ntlm import HttpNtlmAuth
from requests.auth import HTTPBasicAuth from requests.auth import HTTPBasicAuth
from requests.auth import HTTPDigestAuth from requests.auth import HTTPDigestAuth
from requests.sessions import Session from requests.sessions import Session
from .general import debug_dict
#### ####
# Auth URL session # Auth URL session
@ -40,7 +39,8 @@ def call_url(
authtype: str = "basic", authtype: str = "basic",
user: str = "", user: str = "",
passw: str = "", passw: str = "",
timeout: int = 1 timeout: int = 1,
print_errors: bool = False
) -> str: ) -> str:
""" """
Call URL. Uses request module to get url contents. Call URL. Uses request module to get url contents.
@ -55,6 +55,8 @@ def call_url(
Returns: Returns:
- str: call output - str: call output
""" """
from .output import print_stderr
# using with so we make sure the session is closed even when exceptions are encountered # using with so we make sure the session is closed even when exceptions are encountered
with Session() as session: with Session() as session:
if authtype != None: if authtype != None:
@ -65,8 +67,10 @@ def call_url(
try: try:
output = session.get(url, timeout=timeout, verify=False) output = session.get(url, timeout=timeout, verify=False)
except ValueError: except ValueError:
output = "Error: URL format not valid (example http://myserver/page.php)" if print_errors:
print_stderr("Error: URL format not valid (example http://myserver/page.php)")
except Exception as e: except Exception as e:
output = f"{type(e).__name__}:\t{str(e)}" if print_errors:
print_stderr(f"{type(e).__name__}:\t{str(e)}")
return output return output

View File

@ -1,5 +1,3 @@
from .general import debug_dict
#### ####
# Init module template # Init module template
######################################################################################### #########################################################################################
@ -79,6 +77,8 @@ def print_module(
- Module "value" field accepts str type or [list] for datalists. - Module "value" field accepts str type or [list] for datalists.
- Use print_flag to show modules' XML in STDOUT. - Use print_flag to show modules' XML in STDOUT.
""" """
from .output import print_stdout
module_xml = "" module_xml = ""
if module is not None: if module is not None:
@ -193,7 +193,7 @@ def print_module(
module_xml += "</module>\n" module_xml += "</module>\n"
if print_flag: if print_flag:
print(module_xml) print_stdout(module_xml)
return module_xml return module_xml
@ -234,6 +234,8 @@ def print_log_module(
- Module "value" field accepts str type. - Module "value" field accepts str type.
- Use not_print_flag to avoid printing the XML (only populates variables). - Use not_print_flag to avoid printing the XML (only populates variables).
""" """
from .output import print_stdout
module_xml = "" module_xml = ""
if module is not None: if module is not None:
@ -246,6 +248,6 @@ def print_log_module(
module_xml += "</log_module>\n" module_xml += "</log_module>\n"
if print_flag: if print_flag:
print(module_xml) print_stdout(module_xml)
return module_xml return module_xml

View File

@ -0,0 +1,91 @@
import sys
import os
import json
####
# Prints message in stdout
#########################################################################################
def print_stdout(
message: str = ""
):
"""
Prints message in stdout
"""
print(message)
####
# Prints message in stderr
#########################################################################################
def print_stderr(
message: str = ""
):
"""
Prints message in stderr
"""
print(message, file=sys.stderr)
####
# Prints dictionary in formatted json string.
#########################################################################################
def print_debug(
var = "",
print_errors: bool = False
):
"""
Prints any list, dict, string, float or integer as a json
"""
try:
debug_json = json.dumps(var, indent=4)
print_stdout(debug_json)
except json.JSONDecodeError as e:
if print_errors:
print_stderr(f"debug_dict: Failed to dump. Error: {e}")
except Exception as e:
if print_errors:
print_stderr(f"debug_dict: Unexpected error: {e}")
####
# Add new line to log file
#########################################################################################
def logger(
log_file: str = "",
message: str = "",
log_level: str = "",
add_date: bool = True,
print_errors: bool = False
) -> bool:
'''
Add new line to log file
'''
from .general import now
try:
if not os.path.exists(log_file):
with open(log_file, 'w') as file:
pass # Creates an empty file
elif not os.access(log_file, os.W_OK):
if print_errors:
print_stderr(f"Log file '{log_file}' is not writable.")
return False
with open(log_file, 'a') as file:
final_message = ""
if add_date:
final_message += now() + " "
if log_level != "":
final_message += "[" + log_level + "] "
final_message += message + "\n"
file.write(final_message)
return True
except Exception as e:
if print_errors:
print_stderr(f"An error occurred while appending to the log: {e}")
return False

View File

@ -2,7 +2,6 @@ import sys
from queue import Queue from queue import Queue
from threading import Thread from threading import Thread
from multiprocessing import Pool, Manager from multiprocessing import Pool, Manager
from .general import debug_dict
#### ####
# Define multi-processing internal global variables. # Define multi-processing internal global variables.
@ -42,6 +41,7 @@ def run_threads(
""" """
Run a given function for given items list in a given number of threads Run a given function for given items list in a given number of threads
""" """
from .output import print_stderr
# Assign threads # Assign threads
threads = max_threads threads = max_threads
@ -87,7 +87,7 @@ def run_threads(
if print_errors: if print_errors:
for error in errors: for error in errors:
print(error,file=sys.stderr) print_stderr(str(error))
if len(errors) > 0: if len(errors) > 0:
return False return False
@ -96,7 +96,7 @@ def run_threads(
except Exception as e: except Exception as e:
if print_errors: if print_errors:
print("Error while running threads: "+str(e)+"\n",file=sys.stderr) print_stderr("Error while running threads: "+str(e))
return False return False
#### ####
@ -170,6 +170,7 @@ def run_processes(
""" """
Run a given function for given items list in a given number of processes Run a given function for given items list in a given number of processes
""" """
from .output import print_stderr
# Assign processes # Assign processes
processes = max_processes processes = max_processes
@ -187,7 +188,7 @@ def run_processes(
result = True result = True
except Exception as error: except Exception as error:
if print_errors: if print_errors:
print(error,file=sys.stderr) print_stderr(str(error))
result = False result = False
return result return result

View File

@ -4,7 +4,6 @@ import shutil
import subprocess import subprocess
import os import os
import sys import sys
from .general import debug_dict,generate_md5,set_dict_key_value
#### ####
# Define global variables dict, used in functions as default values. # Define global variables dict, used in functions as default values.
@ -35,6 +34,8 @@ def set_global_variable(
variable_name (str): Name of the variable to set. variable_name (str): Name of the variable to set.
value (any): Value to assign to the variable. value (any): Value to assign to the variable.
""" """
from .general import set_dict_key_value
set_dict_key_value(GLOBAL_VARIABLES, variable_name, value) set_dict_key_value(GLOBAL_VARIABLES, variable_name, value)
#### ####
@ -58,6 +59,7 @@ def tentacle_xml(
Returns True for OK and False for errors. Returns True for OK and False for errors.
""" """
from .output import print_stderr
if data_file is not None : if data_file is not None :
@ -70,7 +72,7 @@ def tentacle_xml(
if tentacle_ops['address'] is None : if tentacle_ops['address'] is None :
if print_errors: if print_errors:
sys.stderr.write("Tentacle error: No address defined") print_stderr("Tentacle error: No address defined")
return False return False
try : try :
@ -79,12 +81,12 @@ def tentacle_xml(
data.close() data.close()
except Exception as e : except Exception as e :
if print_errors: if print_errors:
sys.stderr.write(f"Tentacle error: {type(e).__name__} {e}") print_stderr(f"Tentacle error: {type(e).__name__} {e}")
return False return False
tentacle_cmd = f"{tentacle_path} -v -a {tentacle_ops['address']} -p {tentacle_ops['port']} {tentacle_ops['extra_opts']} {data_file.strip()}" tentacle_cmd = f"{tentacle_path} -v -a {tentacle_ops['address']} -p {tentacle_ops['port']} {tentacle_ops['extra_opts']} {data_file.strip()}"
tentacle_exe=Popen(tentacle_cmd, stdout=subprocess.PIPE,stderr=subprocess.PIPE, shell=True) tentacle_exe=subprocess.Popen(tentacle_cmd, stdout=subprocess.PIPE,stderr=subprocess.PIPE, shell=True)
rc=tentacle_exe.wait() rc=tentacle_exe.wait()
if debug == 0 : if debug == 0 :
@ -94,12 +96,12 @@ def tentacle_xml(
if print_errors: if print_errors:
stderr = tentacle_exe.stderr.read().decode() stderr = tentacle_exe.stderr.read().decode()
msg="Tentacle error:" + str(stderr) msg="Tentacle error:" + str(stderr)
print(str(datetime.today().strftime('%Y-%m-%d %H:%M')) + msg , file=sys.stderr) print_stderr(str(datetime.today().strftime('%Y-%m-%d %H:%M')) + msg)
return False return False
else: else:
if print_errors: if print_errors:
sys.stderr.write("Tentacle error: file path is required.") print_stderr("Tentacle error: file path is required.")
return False return False
#### ####
@ -141,7 +143,8 @@ def transfer_xml(
def write_xml( def write_xml(
xml: str = "", xml: str = "",
agent_name: str = "", agent_name: str = "",
data_dir: str = GLOBAL_VARIABLES['temporal'] data_dir: str = GLOBAL_VARIABLES['temporal'],
print_errors: bool = False
) -> str: ) -> str:
""" """
Creates a agent .data file in the specified data_dir folder Creates a agent .data file in the specified data_dir folder
@ -150,6 +153,9 @@ def write_xml(
- agent_name (str): agent name for the xml and file name. - agent_name (str): agent name for the xml and file name.
- data_dir (str): folder in which the file will be created. - data_dir (str): folder in which the file will be created.
""" """
from .general import generate_md5
from .output import print_stderr
Utime = datetime.now().strftime('%s') Utime = datetime.now().strftime('%s')
agent_name_md5 = generate_md5(agent_name) agent_name_md5 = generate_md5(agent_name)
data_file = "%s/%s.%s.data" %(str(data_dir),agent_name_md5,str(Utime)) data_file = "%s/%s.%s.data" %(str(data_dir),agent_name_md5,str(Utime))
@ -158,8 +164,10 @@ def write_xml(
with open(data_file, 'x') as data: with open(data_file, 'x') as data:
data.write(xml) data.write(xml)
except OSError as o: except OSError as o:
print(f"ERROR - Could not write file: {o}, please check directory permissions", file=sys.stderr) if print_errors:
print_stderr(f"ERROR - Could not write file: {o}, please check directory permissions")
except Exception as e: except Exception as e:
print(f"{type(e).__name__}: {e}", file=sys.stderr) if print_errors:
print_stderr(f"{type(e).__name__}: {e}")
return data_file return data_file