class GatherFile(object):
"""GatherFile class."""
def __init__(self, debug=False, path=None):
"""
Initialize GatherFile.
Args:
debug (bool): Log on terminal or not
path (str): Path of the directory to scan files for
Raises:
None
Returns:
None
"""
# Initialize logger
self.logger = AntiVirusLogger(
__name__,
debug=debug
)
# Initialize path of directory to look for
self._PATH = path
def scan_dir(self):
"""
Scan directory to get the list of files.
Args:
None
Raises:
None
Returns:
found_files (list): List of files after scanning
"""
found_files = [] # Initialize empty list of found files
try:
# Iterate through the directory
for root, _, files in os.walk(self._PATH):
for file in files:
found_files.append(os.path.join(root, file))
except Exception as e:
self.logger.log(
"Error occurred: " + str(e),
logtype="error"
)
# Return the list of found files
return found_files
class UpdateHash(object):
"""UpdateHash class."""
def __init__(self, debug=False, config_path=None):
"""
Initialize UpdateHash.
Args:
debug (bool): Log on terminal or not
config_path (str): Configuration JSON file path
Raises:
None
Returns:
None
"""
# Initialize logger
self.logger = AntiVirusLogger(__name__, debug=debug)
if config_path is not None:
self._CONFIG_PATH = config_path
else:
self.logger.log("Configuration file path not found.",
logtype="error")
sys.exit(0)
# Load Configuration
self.config_dict = utils.json_to_dict(self._CONFIG_PATH)
# Categorize OS
self.os_name = utils.categorize_os()
if self.os_name:
try:
# Load hash storage path
self._HASH_STORAGE = self.config_dict[
self.os_name]["update"]["hash"]["storage"]
except KeyError:
self.logger.log("Could not load configuration for: {}".format(
self.os_name),
logtype="error")
sys.exit(0)
# VirusShare Base URL Path
self._HASH_URL = "https://www.virusshare.com/hashes/VirusShare_%05d.md5"
# Max number of Hash files to download
self._MAX = 366
# Create AntiVirus directory, create if not
helper.check_dir(self._HASH_STORAGE)
else:
self.logger.log("Could not determine the OS", logtype="error")
sys.exit(0)
def remove_temp(self):
"""
Remove temporary files generated due to WGET update.
Args:
None
Raises:
None
Returns:
None
"""
files_list = os.listdir(".")
for file in files_list:
# if a temp file is found, remove it
if "VirusShare_" in file and file.endswith(".tmp"):
temp_path = os.path.join(os.getcwd(), file)
self.logger.log("\nRemoving temporary file: " + temp_path,
logtype="info")
os.remove(temp_path)
def update(self):
"""
Start the hash update process.
Args:
None
Raises:
None
Returns:
None
"""
files_list = os.listdir(self._HASH_STORAGE)
last_file_num = 0 # 0 number of hash found
for file in files_list:
if "VirusShare_" in file:
# Get the index number of the hash file
file_num = int(file.split("VirusShare_")[1].split(".md5")[0])
if file_num > last_file_num:
last_file_num = file_num
if last_file_num < self._MAX:
# Start downloading the hashes from the last index number
self.download(last_file_num + 1, self._MAX)
elif last_file_num >= self._MAX:
print("[!] Hash Signatures upto date")
self.logger.log("Hash Signatures upto date", logtype="info")
def download(self, lwr, upr):
"""
Download the hash file within the specified boundaries.
Args:
lwr (int): Starting index of the file to download
upr (int): Last index of the file to download
Raises:
None
Returns:
None
"""
print("\n[!] Updating Hash Signatures")
self.logger.log("Updating Hash Signatures", logtype="info")
for file_num in range(lwr, upr):
dwn_url = self._HASH_URL % file_num
print("\n[!] Downloading {0} of {1} files".format(
file_num, self._MAX))
self.logger.log("Downloading {0} of {1} files".format(
file_num, self._MAX),
logtype="info")
wget.download(dwn_url, out=self._HASH_STORAGE)
print("\n")
self.logger.log("Removing temporary files generated",
logtype="info")
self.remove_temp()
class ScannerEngine(object):
"""ScannerEngine class."""
def __init__(self, debug=False, config_path=None, file_list=None, vt_api_key=None, use_clamav=False, use_yara=False):
"""
Initialize ScannerEngine.
Args:
debug (bool): Log on terminal or not
config_path (str): Configuration JSON file path
vt_api_key (str): VirusTotal API Key
file_list (list): List of files to scan
Raises:
None
Returns:
None
"""
# Initialize logger
self.logger = AntiVirusLogger(
__name__,
debug=debug
)
if config_path is not None:
self._CONFIG_PATH = config_path
else:
self.logger.log(
"Configuration file path not found.",
logtype="error"
)
sys.exit(0)
if file_list:
self.file_list = file_list
else:
# Initialize an empty list
self.file_list = []
# Create HashScanner object
self.hash_scanner = HashScanner(debug=debug,
config_path=self._CONFIG_PATH,
file_list=self.file_list,
vt_api_key=vt_api_key)
# Create YaraScanner object
self.yara_scanner = YaraScanner(debug=debug,
config_path=self._CONFIG_PATH,
file_list=self.file_list,
vt_api_key=vt_api_key)
# Create ClamAVScanner object
self.clamd_scanner = ClamAVScanner(debug=debug,
config_path=self._CONFIG_PATH,
file_list=self.file_list,
vt_api_key=vt_api_key)
# List of process in action
self.process_pool = []
# Store whether to use clamav and yara
self.use_clamav = use_clamav
self.use_yara = use_yara
def start_scanner_engine(self):
"""
Start the scanner engine and stat scanning
the files using three (3) engines in a multi-processing
environment.
1. Hash Scanner Engine
2. Yara Scanner Engine
3. Clam AV Scanner Engine
Args:
None
Raises:
None
Returns:
None
"""
try:
# Create Hash Scanner process
hash_scanner_process = multiprocessing.Process(target=self.hash_scanner.start_scan)
# Add Hash Scanner process to process list
self.process_pool.append(hash_scanner_process)
# Start Hash Scanner process
hash_scanner_process.start()
self.logger.log(
"Hash Scanner engine started",
logtype="info"
)
# Create Yara Scanner process
if self.use_yara:
yara_scanner_process = multiprocessing.Process(target=self.yara_scanner.start_scan)
# Add Yara Scanner process to process list
self.process_pool.append(yara_scanner_process)
# Start Yara Scanner process
yara_scanner_process.start()
self.logger.log(
"Yara Scanner engine started",
logtype="info"
)
# Create Clam AV Scanner process
if self.use_clamav:
clamd_scanner_process = multiprocessing.Process(target=self.clamd_scanner.start_scan)
# Add Clamd AV process to process list
self.process_pool.append(clamd_scanner_process)
# Start Clam AV Scanner process
clamd_scanner_process.start()
self.logger.log(
"Clam AV Scanner engine started",
logtype="info"
)
# Complete the process
for process in self.process_pool:
process.join()
return True
except KeyboardInterrupt:
for process in self.process_pool:
process.terminate()
return True
except Exception as e:
self.logger.log(
"Error occurred: " + str(e),
logtype="error"
)
return True
class MonitorChanges(object):
"""MonitorChanges class."""
def __init__(self,
debug=False,
config_path=None,
min_time=20,
vt_api_key=None,
use_clamav=False,
use_yara=False):
"""
Initialize MonitorChanges class.
Args:
debug (bool): Log on terminal or not
config_path (str): Configuration JSON file path
min_time (int): Minutes before to monitor
vt_api_key (str): Virus Total API Key
Raises:
None
Returns:
None
"""
self.debug = debug
# Initialize logger
self.logger = AntiVirusLogger(__name__, debug=self.debug)
# Initialize time before (minutes) threshold
self._THRESHOLD = int(min_time) * 60
# List of recently modified files
self.modified_files = list()
# Salt to encode file name and their time-stamp
self._SALT = ":@/*&"
# Time to sleep before next monitoring process
self._SLEEP_TIME = 300
self.os_name = utils.categorize_os()
if config_path and self.os_name:
self._CONFIG_PATH = config_path
self.config_data = utils.json_to_dict(self._CONFIG_PATH)
try:
self._PASSWORD_LOG = self.config_data[
self.os_name]["monitor"]["password_log_file"]
except KeyError:
self.logger.log("Could not load password log file",
logtype="error")
sys.exit(0)
except Exception as e:
self.logger.log("Error occurred: " + str(e), logtype="error")
else:
self.logger.log("Configuration file path not found.",
logtype="error")
sys.exit(0)
# Create GatherFile object
self.gather_file_obj = file_gather.GatherFile(path="/")
# List of files recently modified and scanned
self.done_scanning = []
# Virus Total API key
self.vt_api_key = vt_api_key
# UID list
self.verified_uid_list = self.get_initial_uid()
# Store whether to use clamav and yara
self.use_clamav = use_clamav
self.use_yara = use_yara
def get_initial_uid(self):
"""
Return the list of verified UIDs using
the password log files.
Args:
None
Raises:
None
Returns:
temp_list (list): List of extracted UIDs
"""
temp_list = []
log_file_data = utils.open_file(self._PASSWORD_LOG)
for line in log_file_data:
line = line.strip("\n")
data = line.split(":")
uid = data[2]
temp_list.append(int(uid))
return temp_list
def check_uid(self, file):
"""
Check whether the file UID is within the
verified UID list.
Args:
file (str): Path of the file
Raises:
None
Returns:
None
"""
# Get the file UID
uid = os.stat(file).st_uid
# Check if within the list
if uid not in self.verified_uid_list:
return True
def check_file(self, file):
"""
Check whether the file has been recently modified or
has been recently created.
Args:
file (str): Path of the file
Raises:
None
Returns:
None
"""
try:
if self.check_uid(file):
self.logger.log("File: {} UID not valid".format(file),
logtype="warning")
current_time = int(time.time())
modified_time = int(os.path.getmtime(file))
time_diff = int(current_time - modified_time)
salted_file_name = str(file) + self._SALT + str(modified_time)
if time_diff < self._THRESHOLD:
if salted_file_name not in self.done_scanning:
self.modified_files.append(salted_file_name)
self.logger.log(
"File: {} recently modified or created".format(file),
logtype="warning")
except FileNotFoundError:
pass
except Exception as e:
self.logger.log("Error occurred: " + str(e), logtype="error")
def monitor(self):
"""
Monitor for file changes (modification) and new creation.
Args:
None
Raises:
None
Returns:
None
"""
try:
self.logger.log("Monitoring files", logtype="info")
# Gather list of files to monitor
files_list = self.gather_file_obj.scan_dir()
# Check the files
for file in files_list:
self.check_file(file)
# Extract list of files from the salted file list
extracted_list = [
file.split(self._SALT)[0] for file in self.modified_files
]
# Create ScannerEngine object
self.scanner_engine_obj = ScannerEngine(
debug=self.debug,
config_path=self._CONFIG_PATH,
file_list=extracted_list,
vt_api_key=self.vt_api_key,
use_clamav=self.use_clamav,
use_yara=self.use_yara)
# Extend list of scanned files
self.done_scanning.extend(self.modified_files)
# Empty list of modified files for the next round of monitoring
self.modified_files.clear()
# If scanning is complete, go for another round of monitoring
if self.scanner_engine_obj.start_scanner_engine():
# Delete the old ScannerEngine object
del self.scanner_engine_obj
# Rest for specific time to reduce process overload
time.sleep(self._SLEEP_TIME)
# Start the monitoring process again
self.monitor()
except KeyboardInterrupt:
self.logger.log("Keyboard Interrupt detected, quitting monitoring",
logtype="info")
except Exception as e:
self.logger.log("Error occurred: " + str(e), logtype="error")
class SecureTeaAntiVirus(object):
"""SecureTeaAntiVirus class."""
def __init__(self, debug=False, cred=None):
"""
Initialize SecureTeaAntiVirus.
Args:
debug (bool): Log on terminal or not
cred (dict): SecureTea AntiVirus credentials
Raises:
None
Returns:
None
"""
# Initialize logger
self.logger = AntiVirusLogger(__name__, debug=debug)
if not utils.check_root():
self.logger.log("Please run as root exiting.", logtype="error")
sys.exit(0)
if cred is not None:
self.cred = cred
else:
self.logger.log("SecureTea AntiVirus credentials not found.",
logtype="error")
sys.exit(0)
# JSON configuration file path
self._CONFIG_PATH = "securetea/lib/antivirus/config/config.json"
# Initialize required parameters from the credentials passed
self.vt_api_key = self.cred["virustotal-api-key"]
self.update = int(self.cred["update"])
self.monitor_changes = int(self.cred["monitor-file-changes"])
self.monitor_usb = int(self.cred["monitor-usb"])
self.custom_scan = self.cred["custom-scan"]
if self.custom_scan == "":
self.custom_scan = None
self.auto_delete = int(self.cred["auto-delete"])
# Create CoreEngine object
self.core_engine_obj = core_engine.CoreEngine(
debug=debug,
config_path=self._CONFIG_PATH,
vt_api_key=self.vt_api_key,
monitor_changes=self.monitor_changes,
monitor_usb=self.monitor_usb,
update=self.update,
custom_scan=self.custom_scan,
auto_delete=self.auto_delete)
def start(self):
"""
Start AntiVirus core engine.
Args:
None
Raises:
None
Returns:
None
"""
try:
# Start the core engine
self.core_engine_obj.start_engine()
except Exception as e:
self.logger.log("Error occurred: " + str(e), logtype="error")
class USBMonitor(object):
"""USBMonitor class."""
def __init__(self, debug=False, config_path=None, vt_api_key=None):
"""
Initialize USBMonitor.
Args:
debug (bool): Log on terminal or not
config_path (str): Configuration JSON file path
vt_api_key (str): Virus Total API Key
Raises:
None
Returns:
None
"""
self.debug = debug
# Initialize logger
self.logger = AntiVirusLogger(__name__, debug=self.debug)
# Initialize configuration path
if config_path:
self._CONFIG_PATH = config_path
else:
self.logger.log("Configuration file path not found.",
logtype="error")
sys.exit(0)
# Initialize Virus Total API key
self.vt_api_key = vt_api_key
# Create Pyudev Context
self.context = Context()
self.monitor = Monitor.from_netlink(self.context)
# Monitor only USB devices
self.monitor.filter_by(subsystem="usb")
# List of disks
self.disk_list = list()
# Regex to extract names of available disks
self._DISK_REGEX = r"Disk\s([^:]*)"
# Regex to extract disk block data
self._BLOCK_REGEX = "([^/]*)(.*)"
# Command to list disks (Linux)
self.list_disk_command = "fdisk -l"
# Command to list blocks (Linux)
self.list_block_command = "lsblk"
# USB path not found
self.path_found = False
# Logged on screen or not
self.logged = False
# Load initial list of disks
self.create_initial_disk_list()
def create_initial_disk_list(self):
"""
Create list of all disks available at start of
monitoring.
Args:
None
Raises:
None
Returns:
None
"""
output, error = utils.excecute_command(self.list_disk_command)
if output:
output = output.split("\n")
for line in output:
line = line.strip("\n")
found = re.findall(self._DISK_REGEX, line)
if found:
found = found[0].strip(" ")
disk_name = found.split("/")[-1]
if disk_name not in self.disk_list:
self.disk_list.append(disk_name)
elif error:
self.logger.log("Error occurred: " + str(error), logtype="error")
def get_device_path(self):
"""
Get path of the USB device connected.
Args:
None
Raises:
None
Returns:
path (str): Path of the USB device connected
"""
output, error = utils.excecute_command(self.list_disk_command)
if output:
output = output.split("\n")
for line in output:
line = line.strip("\n")
found = re.findall(self._DISK_REGEX, line)
if found:
found = found[0].strip(" ")
disk_name = found.split("/")[-1]
if disk_name not in self.disk_list:
return self.get_block_path(disk_name)
elif error:
self.logger.log("Error occurred: " + str(error), logtype="error")
return
def get_block_path(self, disk_name):
"""
Get block specific path of the USB device.
Args:
disk_name (str): Name of the disk
Raises:
None
Returns:
path (str): Block specific USB device path
"""
str_match = disk_name + self._BLOCK_REGEX
output, error = utils.excecute_command(self.list_block_command)
if output:
output = output.split("\n")
for line in output:
line = line.strip("\n")
found = re.findall(str_match, line)
if found:
path = found[0][1]
if path:
self.path_found = True
return path
elif error:
self.logger.log("Error occurred: " + str(error), logtype="error")
return
def monitor_usb_device(self):
"""
Start monitoring USB devices and scan them.
Args:
None
Raises:
None
Returns:
None
"""
try:
for device in iter(self.monitor.poll, None):
if device.action == "add":
self.path_found = False
self.logged = False
while not self.path_found:
if not self.logged:
self.logger.log("USB Device connected")
self.logged = True
path_name = self.get_device_path()
if path_name:
self.logger.log(
"Scanning USB device, path: {0}".format(
path_name),
logtype="info")
# Create GatherFile object
gather_file_obj = file_gather.GatherFile(
debug=self.debug, path=path_name)
# Get list of files
file_list = gather_file_obj.scan_dir()
# Create ScannerEngine object
self.scanner_engine_obj = scanner_engine.ScannerEngine(
debug=self.debug,
config_path=self._CONFIG_PATH,
vt_api_key=self.vt_api_key,
file_list=file_list)
# Start scanning the files
self.scanner_engine_obj.start_scanner_engine()
if device.action == "remove":
self.logger.log("USB device removed", logtype="info")
except KeyboardInterrupt:
self.logger.log("Exiting USB Monitor", logtype="info")
class Cleaner(object):
"""Cleaner class."""
def __init__(self, debug=False, config_path=None):
"""
Initialize Cleaner.
Args:
debug (bool): Log on terminal or not
config_path (str): Configuration JSON file path
Raises:
None
Returns:
None
"""
# Initialize logger
self.logger = AntiVirusLogger(
__name__,
debug=debug
)
if config_path:
self._CONFIG_PATH = config_path
else:
self.logger.log(
"Configuration file not found",
logtype="error"
)
sys.exit(0)
# Load Configuration
self.config_dict = utils.json_to_dict(self._CONFIG_PATH)
# Categorize OS
self.os_name = utils.categorize_os()
if self.os_name:
# Load malicious-file log path
self._MAL_FILE_PATH = self.config_dict[self.os_name]["scanner"]["malicious_file_log_path"]
def clean(self):
"""
Clean the found malicious files manually.
Args:
None
Raises:
None
Returns:
None
"""
# List of malicious files detected
print("\n[!] List of possible malicious files found: ")
try:
self.malicious_file_list = [file_path.strip("\n") \
for file_path in utils.open_file(self._MAL_FILE_PATH)]
except FileNotFoundError:
# Initialize empty list
self.malicious_file_list = list()
for index, mal_file in enumerate(self.malicious_file_list):
print("[{0}] {1}".format(index, mal_file))
to_clean = input("Enter the index of the file to delete ('e' to exit): ")
if to_clean == "e":
return
else:
to_clean = int(to_clean)
file_path = self.malicious_file_list[to_clean]
print("Removing: ", file_path)
try:
os.remove(file_path)
self.malicious_file_list.remove(file_path)
except FileNotFoundError:
pass
except Exception as e:
self.logger.log(
"Error occurred: " + str(e),
logtype="error"
)
with open(self._MAL_FILE_PATH, "w") as wf:
for mal_file in self.malicious_file_list:
wf.write(mal_file)
self.clean()
def auto_delete(self):
"""
Auto-delete (clean) the found malicious files.
Args:
None
Raises:
None
Returns:
None
"""
self.logger.log(
"Auto-cleaning the malicious files",
logtype="info"
)
print("[!] Auto-cleaning the malicious files")
try:
self.malicious_file_list = [file_path.strip("\n") \
for file_path in utils.open_file(self._MAL_FILE_PATH)]
except FileNotFoundError:
# Initialize empty list
self.malicious_file_list = list()
for mal_file in self.malicious_file_list:
try:
os.remove(mal_file)
self.logger.log(
"Removing: " + mal_file,
logtype="info"
)
self.malicious_file_list.remove(mal_file)
except FileNotFoundError:
pass
except Exception as e:
self.logger.log(
"Error occurred: " + str(e),
logtype="error"
)
with open(self._MAL_FILE_PATH, "w") as wf:
for mal_file in self.malicious_file_list:
wf.write(mal_file)
class CoreEngine(object):
"""CoreEngine class."""
def __init__(self,
debug=False,
config_path="securetea/lib/antivirus/config/config.json",
vt_api_key=None,
use_clamav=False,
use_yara=False,
monitor_changes=1,
monitor_usb=1,
update=1,
custom_scan=None,
auto_delete=0):
"""
Initialize CoreEngine.
Args:
debug (bool): Log on terminal or not
config_path (str): Configuration JSON file path
vt_api_key (str): VirusTotal API Key
monitor_changes (int): Monitor changes (1) or not (0)
monitor_usb (int): Monitor USB (1) or not (0)
Raises:
None
Returns:
None
"""
self.debug = debug
# Initialize logger
self.logger = AntiVirusLogger(
__name__,
debug=debug
)
if config_path is not None:
self._CONFIG_PATH = config_path
else:
self.logger.log(
"Configuration file path not found.",
logtype="error"
)
sys.exit(0)
# Initialize variables
self.update = int(update)
self.monitor_changes = int(monitor_changes)
self.monitor_usb = int(monitor_usb)
self.auto_delete = int(auto_delete)
if custom_scan:
path = custom_scan
else:
path = "~/"
# Create GatherFile object
self.gather_file_obj = GatherFile(path=path)
# Get list of files
self.file_list = self.gather_file_obj.scan_dir()
# Initialize list of process
self.process_pool = []
if vt_api_key and vt_api_key != "XXXX":
self.vt_api_key = vt_api_key
else:
self.vt_api_key = None
self.use_clamav = use_clamav
self.use_yara = use_yara
# Create ScannerEngine object
self.scanner_engine_obj = ScannerEngine(debug=debug,
config_path=self._CONFIG_PATH,
file_list=self.file_list,
vt_api_key=self.vt_api_key,
use_clamav=self.use_clamav,
use_yara=self.use_yara)
# Create MonitorEngine object
self.monitor_engine_obj = MonitorEngine(debug=debug,
config_path=self._CONFIG_PATH,
monitor_changes=self.monitor_changes,
monitor_usb=self.monitor_usb,
vt_api_key=self.vt_api_key,
use_clamav=self.use_clamav,
use_yara=self.use_yara)
# Create Cleaner object
self.cleaner_obj = Cleaner(debug=debug, config_path=self._CONFIG_PATH)
def start_update(self):
"""
Start the update process.
Args:
None
Raises:
None
Returns:
None
"""
# If update is specified
if self.update:
print("[!] Press CTRL+C to skip updates")
# Update Hash
try:
# Create UpdateHash object
self.update_hash_obj = UpdateHash(debug=self.debug,
config_path=self._CONFIG_PATH)
# Start / resume update
self.update_hash_obj.update()
except KeyboardInterrupt:
self.logger.log(
"Skipping Hash update",
logtype="info"
)
print("[!] Skipping Hash update")
self.update_hash_obj.remove_temp()
except Exception as e:
self.logger.log(
"Error occurred: " + str(e),
logtype="error"
)
# Update Yara
try:
# Create UpdateYara object
self.yara_obj = UpdateYara(debug=self.debug,
config_path=self._CONFIG_PATH)
# Start / resume object
self.yara_obj.update()
except KeyboardInterrupt:
self.logger.log(
"Skipping Yara update",
logtype="info"
)
print("[!] Skipping Yara update")
except Exception as e:
self.logger.log(
"Error occurred: " + str(e),
logtype="error"
)
def start_core_process(self):
"""
Start core process (scanner & monitor engine).
Args:
None
Raises:
None
Returns:
None
"""
# Create scanner engine process
scanner_engine_process = multiprocessing.Process(target=self.scanner_engine_obj.start_scanner_engine)
# Create monitor engine process
monitor_engine_process = multiprocessing.Process(target=self.monitor_engine_obj.start_monitor_engine)
# Add scanner process to process pool
self.process_pool.append(scanner_engine_process)
# Add monitor process to process pool
self.process_pool.append(monitor_engine_process)
# Start all the process
for process in self.process_pool:
process.start()
# Complete (join) all the process
for process in self.process_pool:
process.join()
def start_engine(self):
"""
Start core engine of AntiVirus.
Args:
None
Raises:
None
Returns:
None
"""
if self.update: # if update specified
self.start_update() # start the update process
try:
# Start the core process
self.start_core_process()
except KeyboardInterrupt:
for process in self.process_pool:
process.terminate()
except Exception as e:
self.logger.log(
"Error occurred: " + str(e),
logtype="error"
)
# After completing all the process
self.logger.log(
"Collecting found malicious files",
logtype="info"
)
print("[!] Collecting found malicious files, please wait...")
# Sleep for 10 seconds to reset process
time.sleep(10)
# Clear screen
print(chr(27) + "[2J")
# Run the cleaner
if self.auto_delete:
# Auto delete all
self.cleaner_obj.auto_delete()
else:
# Manually delete selected
self.cleaner_obj.clean()
class MonitorEngine(object):
"""
MonitorEngine class.
"""
def __init__(self,
debug=False,
config_path=None,
vt_api_key=None,
monitor_changes=1,
monitor_usb=1):
"""
Initialize MonitorEngine.
Args:
debug (bool): Log on terminal or not
config_path (str): Configuration JSON file path
vt_api_key (str): VirusTotal API Key
monitor_changes (int): Monitor changes (1) or not (0)
monitor_usb (int): Monitor USB (1) or not (0)
Raises:
None
Returns:
None
"""
self.debug = debug
# Initialize logger
self.logger = AntiVirusLogger(__name__, debug=self.debug)
if config_path:
self._CONFIG_PATH = config_path
else:
self.logger.log("Configuration file not found", logtype="error")
sys.exit(0)
# Load Configuration
self.config_dict = utils.json_to_dict(self._CONFIG_PATH)
# Categorize OS
self.os_name = utils.categorize_os()
if self.os_name:
# Load malicious-file log path
self.changes_min_time = int(
self.config_dict[self.os_name]["monitor"]["threshold_min"])
self.monitor_changes = int(monitor_changes)
self.monitor_usb = int(monitor_usb)
# Create a pool of process
self.process_pool = []
# Initialize VirusTotal API key
self.vt_api_key = vt_api_key
def kill_process(self):
"""
Kill running process.
Args:
None
Raises:
None
Returns:
None
"""
for process in self.process_pool:
process.terminate()
def create_process(self):
"""
Create specific process depending on the choice
of the user.
Args:
None
Raises:
None
Returns:
None
"""
if self.monitor_changes:
# Create MonitorChanges object
self.monitor_changes_obj = MonitorChanges(
debug=self.debug,
config_path=self._CONFIG_PATH,
min_time=self.changes_min_time,
vt_api_key=self.vt_api_key)
monitor_changes_process = multiprocessing.Process(
target=self.monitor_changes_obj.monitor)
# Add to process pool
self.process_pool.append(monitor_changes_process)
if self.monitor_usb:
# Create USBMonitor object
self.monitor_usb_obj = USBMonitor(debug=self.debug,
config_path=self._CONFIG_PATH,
vt_api_key=self.vt_api_key)
monitor_usb_process = multiprocessing.Process(
target=self.monitor_usb_obj.monitor_usb_device)
# Add to process pool
self.process_pool.append(monitor_usb_process)
def start_monitor_engine(self):
"""
Start the monitor engine.
Args:
None
Raises:
None
Returns:
None
"""
# Create process based on user choice
self.create_process()
try:
if self.process_pool:
for process in self.process_pool:
process.start()
for process in self.process_pool:
process.join()
except KeyboardInterrupt:
self.logger.log(
"KeyboardInterrupt detected, quitting monitor engine",
logtype="info")
# Kill running process
self.kill_process()
except Exception as e:
self.logger.log("Error occurred: " + str(e), logtype="error")
# Kill running process
self.kill_process()
class VirusTotal(object):
"""VirusTotal class."""
def __init__(self, debug=False, api_key=None):
"""
Initialize VirusTotal.
Args:
debug (bool): Log on terminal or not
api_key (str): VirusTotal API Key
Raises:
None
Returns:
None
"""
# Initialize logger
self.logger = AntiVirusLogger(__name__, debug=debug)
# VirusTotal End-point API URL
self._API_URL = "https://www.virustotal.com/vtapi/v2/file/report"
if api_key:
self.api_key = api_key
# Error to Reason mapping (as per VirusTotal website)
self.error_status_code_map = {
"203":
"Request rate limit exceeded. You are making more requests "
"than allowed. You have exceeded one of your quotas (minute, "
"daily or monthly). Daily quotas are reset every day at 00:00 UTC.",
"400":
"Bad request. Your request was somehow incorrect. "
"This can be caused by missing arguments or arguments "
"with wrong values.",
"403":
"Forbidden. You don't have enough privileges to make the "
"request. You may be doing a request without providing an "
"API key or you may be making a request to a Private API "
"without having the appropriate privileges."
}
def check_hash(self, hash_value, file_path):
"""
Check the MD5 Hash value to detect any possible
malicious file.
Args:
hash_value (str): MD5 hash value of the file to check
file_path (str): Path of the file
Raises:
None
Returns:
None
"""
# Get the name of the file
file_name = file_path.split("/")[-1]
# Set up parameters using the VirusTotal API format
params = {"apikey": self.api_key, "resource": hash_value}
resp = requests.get(self._API_URL, params=params)
status = resp.status_code
if status == 200: # if a success
resp = resp.json()
positives = resp["positives"]
if int(
positives
) >= 1: # more than one antivirus detected file as suspicious
self.logger.log(
"File: {0} found suspicious in VirusTotal SandBox test".
format(file_name),
logtype="warning")
return True
else:
self.logger.log(
"File: {0} not found suspicious in VirusTotal SandBox test"
.format(file_name),
logtype="info")
return False
elif self.error_status_code_map.get(str(status)):
self.logger.log(self.error_status_code_map[str(status)],
logtype="error")
else:
self.logger.log(
"VirusTotal API: Could not fetch information, error code: {0}".
format(status),
logtype="error")
class Scanner(object):
"""Scanner class."""
def __init__(self,
debug=False,
config_path=None,
file_list=None,
vt_api_key=None):
"""
Initialize Scanner.
Args:
debug (bool): Log on terminal or not
config_path (str): Configuration JSON file path
file_list (list): List of files to scan
vt_api_key (str): Virus Total API Key
Raises:
None
Returns:
None
"""
# Initialize logger
self.logger = AntiVirusLogger(__name__, debug=debug)
if config_path is not None:
self._CONFIG_PATH = config_path
else:
self.logger.log("Configuration file path not found.",
logtype="error")
sys.exit(0)
# Load Configuration
self.config_dict = utils.json_to_dict(self._CONFIG_PATH)
# Categorize OS
self.os_name = utils.categorize_os()
if self.os_name:
# Load malicious-file log path
self._MAL_FILE_PATH = self.config_dict[
self.os_name]["scanner"]["malicious_file_log_path"]
if file_list is not None:
self.file_list = file_list
else:
self.file_list = []
# List of malicious files detected
try:
self.malicious_file_list = [file_path.strip("\n") \
for file_path in utils.open_file(self._MAL_FILE_PATH)]
except FileNotFoundError:
# Initialize empty list
self.malicious_file_list = list()
self.vt_api_key = vt_api_key
if self.vt_api_key and self.vt_api_key != "XXXX":
# If VirusTotal API Key is provided & valid
self.vt_obj = VirusTotal(debug=debug, api_key=self.vt_api_key)
def scan_file(self, file_path):
"""
Scan file.
Args:
file_path (str): Path of the file to scan
Raises:
None
Returns:
None
"""
# Over-write this function to the desired logic.
pass
def start_scan(self):
"""
Start scanning process in a
multi-threaded environment.
Args:
None
Raises:
None
Returns:
None
"""
try:
with ThreadPoolExecutor(max_workers=self._WORKERS) as executor:
executor.map(self.scan_file, self.file_list)
except KeyboardInterrupt:
self.logger.log("Keyboard Interrupt detected, quitting scan",
logtype="info")
except Exception as e:
self.logger.log("Error occurred: " + str(e), logtype="error")
def check_virus_total(self, file_path):
"""
Check the file using Virus Total API Sand Box.
Args:
file_path (str): Path of the file
Raises:
None
Returns:
None
"""
# If VirusTotal API Key is provided
if self.vt_api_key:
# Perform a third confirmation test
self.logger.log(
"Testing malicious file: {0} under VirusTotal Sand Box".format(
file_path),
logtype="info")
file_hash_value = utils.get_md5_hash(file_path=file_path)
if self.vt_obj.check_hash(hash_value=file_hash_value,
file_path=file_path):
self.logger.log(
"File: {0} found malicious in VirusTotal Sand Box Test".
format(file_path),
logtype="warning")
utils.write_data(self._MAL_FILE_PATH, file_path)
else:
self.logger.log(
"File: {0} not found malicious in VirusTotal Sand Box Test"
.format(file_path),
logtype="info")
else: # Skip the third confirmation test
self.logger.log(
"Skipping VirusTotal Sand Box test for possible malicious file: {0}"
.format(file_path),
logtype="info")
utils.write_data(self._MAL_FILE_PATH, file_path)
return
class UpdateYara(object):
"""UpdateYara class."""
def __init__(self, debug=False, config_path=None):
"""
Initialize UpdateYara.
Args:
debug (bool): Log on terminal or not
config_path (str): Configuration JSON file path
Raises:
None
Returns:
None
"""
# Initialize logger
self.logger = AntiVirusLogger(__name__, debug=debug)
if config_path is not None:
self._CONFIG_PATH = config_path
else:
self.logger.log("Configuration file path not found.",
logtype="error")
sys.exit(0)
# Yara Download URL
self._YARA_DW_URL = "https://raw.githubusercontent.com/Yara-Rules/rules/master/malware/"
# Yara GitHub Repo URL to update list of rules
self._YARA_NAMELIST_URL = "https://github.com/Yara-Rules/rules/tree/master/malware"
# Match string to get names of rules
self._YARA_MATCH = "/Yara-Rules/rules/blob/master/malware/"
# Load Configuration
self.config_dict = utils.json_to_dict(self._CONFIG_PATH)
# Categorize OS
self.os_name = utils.categorize_os()
if self.os_name:
try:
# Load Yara storage path
self._YARA_STORAGE = self.config_dict[
self.os_name]["update"]["yara"]["storage"]
except KeyError:
self.logger.log("Could not load configuration for: {}".format(
self.os_name),
logtype="error")
sys.exit(0)
else:
self.logger.log("Could not determine the OS", logtype="error")
sys.exit(0)
# Check whether Yara rules storage directory exists, create one if not
helper.check_dir(self._YARA_STORAGE)
# List of Yara rules to download
self.name_list = []
# List of already downloaded (updated) rules
self.downloaded = self.current_status()
# Set flag to no download required (default)
self.flag = 0
def get_namelist(self):
"""
Collect names of Yara rules from the GitHub repo.
Args:
None
Raises:
None
Returns:
None
"""
response = requests.get(self._YARA_NAMELIST_URL)
soup_obj = BeautifulSoup(response.text, "lxml")
a_tags = soup_obj.find_all("a")
for a_tag in a_tags:
link = a_tag.get("href")
if self._YARA_MATCH in link:
name = link.split(self._YARA_MATCH)[1]
name = name.strip(" ")
if name not in self.name_list:
self.name_list.append(name)
def current_status(self):
"""
Return list of downloaded (updated) Yara rules.
Args:
None
Raises:
None
Returns:
downloaded (list): List of downloaded (updated) Yara rules
"""
downloaded = os.listdir(self._YARA_STORAGE)
return downloaded
def update(self):
"""
Start the Yara rules update process.
Args:
None
Raises:
None
Returns:
None
"""
# Get list of already downloaded Yara rules
self.get_namelist()
for name in self.name_list:
if name not in self.downloaded:
self.flag = 1 # Download in process
print("\n[!] Downloading: ", name)
self.logger.log("Downloading: {}".format(name), logtype="info")
dwn_url = self._YARA_DW_URL + name
wget.download(dwn_url, out=self._YARA_STORAGE)
if self.flag == 0: # No download needed
print("\n[!] Yara rules upto date")
self.logger.log("Yara rules upto date", logtype="info")
