#!/usr/bin/python

import sys
from shutil import copyfile
import filecmp
import os
import subprocess

def updateFile (status, snortRulesFileName):
    print "working with file " + snortRulesFileName
    ruleFile = open(snortRulesFileName, 'r')
    w = open('output.txt', 'w')
    
    if status.lower() == "enable":
        for line in ruleFile:
            # Get the first 7 characters
            chars = line[:7]
            if chars == "# alert":
                    line = line[2:] # Remove first two beginning characters
            # Write the rule to another file
            w.write(line)

    if status.lower() == "disable":
        for line in ruleFile:
            # Add the comment chars
            # Get the first five characters of the line
            chars = line[:5]
            if chars == "alert":
                line = '# ' + line
                # print line
            w.write(line)
    ruleFile.close()
    # Copy the output file to the original file
    #copyfile('output.txt', snortRulesFileName)
    w.close()
    res = subprocess.check_output(["cp", "output.txt", snortRulesFileName])
    for line in res.splitlines():
        print line
    #subprocess.Popen("cp output.txt " + snortRulesFileName)
    # Then delete the output file
    os.remove('output.txt')

try:
    n, status, snortRulesFileName = sys.argv
except ValueError:
    print "Description: This script (un)comments rules in Snort rule files.  " \
          "Modifying of rule lines starts at line 30, and goes till the end of the rule file.  " \
          "Note: To uncomment rules, the rules must start in the following format: \"# \".  " \
          "To comment rules, there must not be any text before the rule, nor spaces.\n" \
          "Arguments: Status[enable/disable] Name of Rule File"
    sys.exit(1)

isaDirectory = os.path.isdir(snortRulesFileName)
isaFile = os.path.isfile(snortRulesFileName)

if(isaDirectory):
    print "I notice this is a directory.  Going to recursively modify the files"
    files = os.listdir(snortRulesFileName)
    for filename in files:
	print "CHecking file: " + "/etc/snort/rules/"+filename
        if os.path.isfile("/etc/snort/rules/"+filename):
	    print "this is a file!!!"
            updateFile(status, snortRulesFileName + filename) 
else:
    print "File detected." + snortRulesFileName
    updateFile(status, snortRulesFileName)

# Default input is to generate backup file
user_response = raw_input("Would you first like to make a backup of the snort rules? [Y/n]: ")
while user_response != "y" and user_response != "" and user_response != "n":
    user_response = raw_input("Invalid user input.  Please try again.\n"
                              "Would you first like to make a backup of the snort rules? [Y/n]: ")

if user_response == "y" or user_response == "":
    # Name of the new backup file we are going to generate
    backupFilename = snortRulesFileName + ".orig"
    # Make a backup of the file with the extension of ".backup"
    print "Making backup file..."
    # Make a copy of the file
    copyfile(snortRulesFileName, backupFilename)
    print "Checking if backup file is identical to original..."
    # Check to see if identical file matches contents of new backup file before proceeding
    result = cmp(snortRulesFileName, backupFilename)
    if result:
        print "Success!  Continuing..."
    else:
        print "[!] Failure.  Exiting."
        sys.exit(1)
else:
    # Don't make a backup of the file
    print "OK, don't say I didn't ask! Skipping backup of file"


#modifiedLines = "{:,}".format(modifiedLines)
#print "Script done executing. " + status + " " + str(modifiedLines) + " rules."