#!/usr/bin/env python
# Author: Justin Konecny

from __builtin__ import file
from _ast import Num
import argparse
from datetime import datetime
import hashlib
import os
import random
import sys


# executes the program
def main():
    # starts timer to track program runtime
    start = datetime.now()
    print("[START]: " + start.__str__())
    
    # creates the command line argument parser and the appropriate required inputs
    parser = argparse.ArgumentParser(description="Anti-virus training software")
    parser.add_argument('dir_malware', metavar='<malware directory>', action='store', help="the directory of training malware")
    parser.add_argument('dir_benign', metavar='<benign directory>', action='store', help="the directory of benign training software")
    parser.add_argument('file_sig_output', metavar='<output signature file>', action='store', help="the output file for the malware signatures")

    # reads and stores the command arguments
    args = parser.parse_args()
    dir_malware = args.dir_malware
    dir_benign = args.dir_benign
    file_output = args.file_sig_output
    
    # outputs the arguments to the console
    print("[Malware Directory]: " + dir_malware)
    print("[Benign Directory]: " + dir_benign)
    print("[Signature File Output]: " + file_output)
    
    try:
        # tries to find the paths to the given directories
        path_malware = os.listdir(dir_malware)
        path_benign = os.listdir(dir_benign)
    except OSError as e: 
        # exits the program if it cannot find the directories
        print("[Directory Error]: " + e.__str__())
        return
    
    # creates the malware signatures
    createSignatures(file_output, path_malware, dir_malware, path_benign, dir_benign, 48)
    
    # prints the program runtime
    print("[RUNTIME]:" + (datetime.now() - start).__str__())
    

# creates the malware signatures given: output file, malware path,
# malware directory, safeware path, safeware directory, signature length (in bytes)
def createSignatures(file_output, path_malware, dir_malware, path_benign, dir_benign, num):
    print(os.linesep + "[[==Creating MALWARE Signatures==]]")
    # creates and opens the signatures file as writable
    file_signatures = open(file_output, 'w+')
    # creates an empty list of signatures to prevent duplicates
    list_sig = []
    # counters for the number of signatures created/failed
    count_success = 0
    count_fail = 0
    # opens every file in the malware directory
    for file in path_malware:
        with open(dir_malware + file) as content:
            content.read(num * 2) # prevents signatures from being the first (num * 2) bytes
            try:
                # finds a unique string in the malware
                string = findString(list_sig, content, path_benign, dir_benign, num)
                
                # writes the string if it is unique to this malware only
                if not (string == True):
                    # writes this string to the signature file
                    file_signatures.write(string + os.linesep)
                    # appends the string to the list of signatures
                    list_sig.append(string)
                
                # increments the success counter
                count_success += 1
                print("[" + str('{0:04}'.format(count_success)) + "][Signature Created For MALWARE]: " + file)
            except RuntimeError as e:
                # increments the failure counter
                count_fail += 1
                print("[NO SIGNATURE CREATED FOR MALWARE]: " + file)
                print("[Error]: " + e.__str__())
    # closes the signature file         
    file_signatures.close()
    # prints results to the console
    print("[#SUCCEEDED]: " + count_success.__str__())
    print("[#FAILED]: " + count_fail.__str__())
    print("[[==Finished Writing MALWARE Signatures==]]")


# finds a unique string in the given file content, relative to the given safeware files, given: list of
# valid signatures, malware content, path to safeware, directory to safeware, and signature length (in bytes)
def findString(list_sig, content, path_benign, dir_benign, num):
    # reads the next num bytes from the given content
    string = content.read(num)
    # checks if the string is already in the list of valid signatures
    if string in list_sig:
        return True # represents a valid signature that has already been written
    else: 
        # checks if this string is a unique/valid signature
        valid = isValidSignature(string, path_benign, dir_benign, num)
        if not valid:
            # recursively executes until a valid string is found
            return findString(list_sig, content, path_benign, dir_benign, num)
        else:
            return string
    

# checks if the given string is a valid signature, relative to the safeware in the given directory
# given signature string, path to safeware, directory to safeware, signature length (in bytes)
def isValidSignature(pot_sig, path_benign, dir_benign, num):
    if os.linesep in pot_sig:
        # invalid signature if it contains a newline character
        return False
    else:
        # checks signature against every piece of safe software
        for file_safe in path_benign:
            with open(dir_benign + file_safe) as content_safe:
                # checks if the safe software contains the string
                content_all = content_safe.read()
                if pot_sig in content_all:
                    return False
        return True


# invokes main() if program is run from the command line
if __name__ == '__main__':
    main()

