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,7 +1,9 @@
from .general import *
from .output import *
from .encryption import *
from .threads import *
from .agents import *
from .modules import *
from .transfer import *
from .discovery import *
from .http import *
from .http import *

View File

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

View File

@ -1,6 +1,5 @@
import sys
import json
from .general import debug_dict
####
# 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
the 'ERROR_LEVEL' as the exit code.
"""
from .output import print_stdout
global ERROR_LEVEL
global SUMMARY
@ -155,5 +155,5 @@ def print_output():
json_string = json.dumps(OUTPUT)
print(json_string)
print_stdout(json_string)
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 json
from datetime import datetime
import hashlib
@ -323,24 +322,6 @@ def safe_output(
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.
#########################################################################################
@ -386,13 +367,15 @@ def generate_md5(
#########################################################################################
def now(
print_flag: int = 0,
utimestamp: int = 0
print_flag: bool = False,
utimestamp: bool = False
) -> str:
"""
Returns time in yyyy/mm/dd HH:MM:SS format by default. Use 1 as an argument
to get epoch time (utimestamp)
"""
from .output import print_stdout
today = datetime.today()
if utimestamp:
@ -401,7 +384,7 @@ def now(
time = today.strftime('%Y/%m/%d %H:%M:%S')
if print_flag:
print(time)
print_stdout(time)
return time
@ -443,6 +426,8 @@ def parse_configuration(
Returns:
- dict: containing all keys and values from file.
"""
from .output import print_stderr
config = {}
try:
@ -456,7 +441,7 @@ def parse_configuration(
config[option.strip()] = value.strip()
except Exception as e:
print (f"{type(e).__name__}: {e}")
print_stderr(f"{type(e).__name__}: {e}")
for option, value in default_values.items():
if option.strip() not in config:
@ -472,7 +457,7 @@ def parse_csv_file(
file: str = "",
separator: str = ';',
count_parameters: int = 0,
debug: bool = False
print_errors: bool = False
) -> 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
- separator (str, optional): Separator for option and value. Defaults to ";".
- 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:
- List: containing a list for of values for each csv line.
"""
from .output import print_stderr
csv_arr = []
try:
@ -498,11 +485,11 @@ def parse_csv_file(
value = line.strip().split(separator)
if len(value) >= count_parameters:
csv_arr.append(value)
elif debug==True:
print(f'Csv line: {line} does not match minimun parameter defined: {count_parameters}',file=sys.stderr)
elif print_errors==True:
print_stderr(f'Csv line: {line} does not match minimun parameter defined: {count_parameters}')
except Exception as e:
print (f"{type(e).__name__}: {e}")
print_stderr(f"{type(e).__name__}: {e}")
return csv_arr

View File

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

View File

@ -1,5 +1,3 @@
from .general import debug_dict
####
# Init module template
#########################################################################################
@ -79,6 +77,8 @@ def print_module(
- Module "value" field accepts str type or [list] for datalists.
- Use print_flag to show modules' XML in STDOUT.
"""
from .output import print_stdout
module_xml = ""
if module is not None:
@ -193,7 +193,7 @@ def print_module(
module_xml += "</module>\n"
if print_flag:
print(module_xml)
print_stdout(module_xml)
return module_xml
@ -234,6 +234,8 @@ def print_log_module(
- Module "value" field accepts str type.
- Use not_print_flag to avoid printing the XML (only populates variables).
"""
from .output import print_stdout
module_xml = ""
if module is not None:
@ -246,6 +248,6 @@ def print_log_module(
module_xml += "</log_module>\n"
if print_flag:
print(module_xml)
print_stdout(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 threading import Thread
from multiprocessing import Pool, Manager
from .general import debug_dict
####
# 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
"""
from .output import print_stderr
# Assign threads
threads = max_threads
@ -87,7 +87,7 @@ def run_threads(
if print_errors:
for error in errors:
print(error,file=sys.stderr)
print_stderr(str(error))
if len(errors) > 0:
return False
@ -96,7 +96,7 @@ def run_threads(
except Exception as e:
if print_errors:
print("Error while running threads: "+str(e)+"\n",file=sys.stderr)
print_stderr("Error while running threads: "+str(e))
return False
####
@ -170,6 +170,7 @@ def run_processes(
"""
Run a given function for given items list in a given number of processes
"""
from .output import print_stderr
# Assign processes
processes = max_processes
@ -187,7 +188,7 @@ def run_processes(
result = True
except Exception as error:
if print_errors:
print(error,file=sys.stderr)
print_stderr(str(error))
result = False
return result

View File

@ -4,7 +4,6 @@ import shutil
import subprocess
import os
import sys
from .general import debug_dict,generate_md5,set_dict_key_value
####
# 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.
value (any): Value to assign to the variable.
"""
from .general import set_dict_key_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.
"""
from .output import print_stderr
if data_file is not None :
@ -70,7 +72,7 @@ def tentacle_xml(
if tentacle_ops['address'] is None :
if print_errors:
sys.stderr.write("Tentacle error: No address defined")
print_stderr("Tentacle error: No address defined")
return False
try :
@ -79,12 +81,12 @@ def tentacle_xml(
data.close()
except Exception as e :
if print_errors:
sys.stderr.write(f"Tentacle error: {type(e).__name__} {e}")
print_stderr(f"Tentacle error: {type(e).__name__} {e}")
return False
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()
if debug == 0 :
@ -94,12 +96,12 @@ def tentacle_xml(
if print_errors:
stderr = tentacle_exe.stderr.read().decode()
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
else:
if print_errors:
sys.stderr.write("Tentacle error: file path is required.")
print_stderr("Tentacle error: file path is required.")
return False
####
@ -141,7 +143,8 @@ def transfer_xml(
def write_xml(
xml: str = "",
agent_name: str = "",
data_dir: str = GLOBAL_VARIABLES['temporal']
data_dir: str = GLOBAL_VARIABLES['temporal'],
print_errors: bool = False
) -> str:
"""
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.
- 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')
agent_name_md5 = generate_md5(agent_name)
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:
data.write(xml)
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:
print(f"{type(e).__name__}: {e}", file=sys.stderr)
if print_errors:
print_stderr(f"{type(e).__name__}: {e}")
return data_file