MalwareSig.py

# This is our code!
# This class will help us interact with the user and our YaraScanner class to scan files and identify 'malware'
# Date: 11/21/19
# Version 1.0
# Authors:
    # Patrick Sacchet
    # Javon Kitson

import sys
import os
import yara

from YaraScanner import *

# Function prints help statement if argument is found
# Input: None
# Returns: None
def print_help():
    print("Here's how you run our code:")
    print("\tpython MalwareSig.py -arguments")
    print("Optional arguments you can add if you'd like:")
    print("\t-h - Print help message\n" +
    "\t-r - Pass rule files (requires malicious file(s) passed in addition)\n" +
    "\t-m - Malicious file(s) to be scanned (requires rule files passed in addition)\n" +
    "If no arguments are passed, program will scan test files with test rule files by default")
    return

# Function prints usage statement since user called program incorrectly
# Input: None
# Returns: None
def print_usage():
    print("Invlalid call of MalwareSig.py, call as follows:")
    print("\tpython MalwareSig.py -arguments")
    print("Optional arguments:")
    print("\t-h - Print help message\n" +
    "\t-r - Pass rule files (requires malicious file(s) passed in addition)\n" +
    "\t-m - Malicious file(s) to be scanned (requires rule files passed in addition)\n" +
    "If no arguments are passed, program will scan test files with test rule files by default")
    return

# Function will parse command line arguments, call functions accordingly dependent on arguments passed
# Input: args - arguments grabbed from the command line
# Returns: YaraScanner object with file paths as attribtues
def parse_commands(args):
    # User passed in help argument, print help message
    if("-h" in args):
        print_help()
    # Otherwise, check for malacious file and rule arguments and check for their files to be passed in as well
    elif(("-r" in args) and ("-m" in args)):
        rule_ind = args.index('-r')
        mal_ind = args.index('-m')
        try:
            # We want to make sure user actually passed in values following command line arguments
            rule_file_loc = args[rule_ind + 1]
            assert(rule_file_loc != "-m")
            mal_file_loc = args[mal_ind + 1]
            assert(mal_file_loc != "-r")
        except Exception as e:
            print("You need to pass in a rule file location and malicious file location, try again")
            exit(0)
        # Create our instance of an YaraScanner object with the file paths
        yara_scanner = YaraScanner(rule_file_loc, mal_file_loc)
        return yara_scanner
    # If user did not pass in file paths for rules and malicious files, just use the one in the test directories
    else:
        print("No parameters or incorrect parameters detected, running with default test directories...")
        yara_scanner = YaraScanner("rules", "malware_files")
        return yara_scanner

# Function will check all hit files and print results accordingly
# Input: hit_files - list of hit files (if any)
# Returns: None (prints results)
def print_results(hit_files):
    if(len(hit_files) < 1):
        print("No files were hit given your Yara rules")
    else:
        for filename in hit_files:
            print("File was hit: " + filename)
    return

# Main functionality
def main():
    # Create a instance of a YaraScanner object with users file paths if passed in
    yara_scanner = parse_commands(sys.argv)
    # Create a rule dictionary so we can compile Yara rules
    yara_rule_dict = yara_scanner.make_dict()
    # Look for hit files using the scanner
    hit_files = yara_scanner.scan_files(yara_rule_dict)
    print_results(hit_files)
    print("For specifics on why these files were hit, check  out " + str(yara_scanner.log))
    return




if __name__ == '__main__':
    main()