update_tw_rules.py

#!/usr/local/bin/python
import requests
import argparse
from argparse import RawTextHelpFormatter
import json
import sys
import time
import os

class Printer():
    """
    Print things to stdout on one line dynamically
    """
    def __init__(self,data):
        sys.stdout.write("\r\x1b[K"+data.__str__())
        sys.stdout.flush()

class AutoVivification(dict):
    """Implementation of perl's autovivification feature."""
    def __getitem__(self, item):
        try:
            return dict.__getitem__(self, item)
        except KeyError:
            value = self[item] = type(self)()
            return value

tw_mapping = {"204208":"2004208", "204209":"2004209", "204210":"2004210", "204211":"2004211",
  "204212":"2004212", "204213":"2004213", "204214":"2004214", "204215":"2004215",
  "204219":"2004219", "204243":"2004243", "204244":"2004244", "204529":"2004529",
  "204530":"2004530", "207369":"2007369", "207377":"2007377", "207368":"2007368",
  "207372":"2007372", "207383":"2007383", "206308":"2006308", "202186":"2002186",
  "202875":"2002875", "203302":"2003302", "204258":"2004258", "204270":"2004270",
  "204277":"2004277", "204279":"2004279", "204644":"2004644", "207809":"2007809",
  "208353":"2008353", "208355":"2008355", "202193":"2002193", "204256":"2004256",
  "204257":"2004257", "204259":"2004259", "204275":"2004275", "204276":"2004276",
  "204297":"2004297", "204471":"2004471", "207810":"2007810", "207947":"2007947",
  "207948":"2007948", "208164":"2008164", "208165":"2008165", "208166":"2008166",
  "203344":"2003344", "209993":"2009993", "204309":"2004309", "204310":"2004310",
  "204311":"2004311", "204312":"2004312", "205567":"2005567", "205568":"2005568",
  "206815":"2006815", "207256":"2007256", "217055":"2017055", "2022":"2000022",
  "202479":"2002479", "202511":"2002511", "20265":"2000265", "202713":"2002713",
  "202722":"2002722", "20278":"2000278", "202829":"2002829", "203012":"2003012",
  "203396":"2003396", "203397":"2003397", "203746":"2003746", "203800":"2003800",
  "20393":"2000393", "203993":"2003993", "204026":"2004026", "20481":"2000481", "20483":"2000483",
  "20498":"2000498", "20506":"2000506", "20508":"2000508", "20512":"2000512", "205169":"2005169",
  "20528":"2000528", "20563":"2000563", "20575":"2000575", "205758":"2005758", "20590":"2000590",
  "206300":"2006300", "206459":"2006459", "20647":"2000647", "206661":"2006661",
  "20677":"2000677", "20693":"2000693", "20699":"2000699", "207012":"2007012", "207112":"2007112",
  "207335":"2007335", "207400":"2007400", "207707":"2007707", "207715":"2007715",
  "207722":"2007722", "207728":"2007728", "208192":"2008192", "208367":"2008367",
  "208882":"2008882", "208956":"2008956", "208983":"2008983", "209054":"2009054",
  "209900":"2009900", "209954":"2009954", "202119":"2002119", "202584":"2002584",
  "202640":"2002640", "202703":"2002703", "202758":"2002758", "202927":"2002927",
  "202959":"2002959", "203059":"2003059", "203319":"2003319", "203342":"2003342",
  "203508":"2003508", "203618":"2003618", "203901":"2003901", "203915":"2003915",
  "203971":"2003971", "203973":"2003973", "203995":"2003995", "204047":"2004047",
  "204103":"2004103", "204294":"2004294", "204307":"2004307", "204495":"2004495",
  "204496":"2004496", "204497":"2004497", "204516":"2004516", "204606":"2004606",
  "204610":"2004610", "204638":"2004638", "204675":"2004675", "204771":"2004771",
  "204785":"2004785", "204786":"2004786", "204838":"2004838", "204931":"2004931",
  "20509":"2000509", "205211":"2005211", "205212":"2005212", "205213":"2005213",
  "205214":"2005214", "205215":"2005215", "205216":"2005216", "205217":"2005217",
  "205218":"2005218", "205219":"2005219", "205220":"2005220", "205221":"2005221",
  "205222":"2005222", "205223":"2005223", "205224":"2005224", "205225":"2005225",
  "205226":"2005226", "205227":"2005227", "205228":"2005228", "205229":"2005229",
  "205262":"2005262", "205329":"2005329", "205368":"2005368", "205369":"2005369",
  "205385":"2005385", "205625":"2005625", "205653":"2005653", "205655":"2005655",
  "205659":"2005659", "205661":"2005661", "205664":"2005664", "205752":"2005752",
  "206013":"2006013", "206244":"2006244", "206358":"2006358", "206360":"2006360",
  "206387":"2006387", "206400":"2006400", "206440":"2006440", "206503":"2006503",
  "206533":"2006533", "207020":"2007020", "207483":"2007483", "207484":"2007484",
  "207485":"2007485", "207486":"2007486", "207783":"2007783", "208031":"2008031",
  "208169":"2008169", "208170":"2008170", "208171":"2008171", "208172":"2008172",
  "208173":"2008173", "208174":"2008174", "208175":"2008175", "208176":"2008176",
  "208177":"2008177", "208178":"2008178", "208179":"2008179", "208181":"2008181",
  "208258":"2008258", "208320":"2008320", "208442":"2008442", "208976":"2008976",
  "209044":"2009044", "209083":"2009083", "209098":"2009098", "209392":"2009392",
  "209497":"2009497", "209510":"2009510", "209918":"2009918", "209920":"2009920",
  "209921":"2009921", "209937":"2009937", "209938":"2009938", "209939":"2009939",
  "209940":"2009940", "209941":"2009941", "209942":"2009942", "209943":"2009943",
  "209944":"2009944", "209945":"2009945", "209946":"2009946", "209947":"2009947",
  "202137":"2002137", "202191":"2002191", "202192":"2002192", "202322":"2002322",
  "202390":"2002390", "202480":"2002480", "202617":"2002617", "202669":"2002669",
  "202680":"2002680", "202762":"2002762", "202767":"2002767", "202788":"2002788",
  "202790":"2002790", "202820":"2002820", "202895":"2002895", "202933":"2002933",
  "203073":"2003073", "203125":"2003125", "203136":"2003136", "203220":"2003220",
  "203273":"2003273", "203281":"2003281", "203330":"2003330", "203335":"2003335",
  "203353":"2003353", "203393":"2003393", "203400":"2003400", "203473":"2003473",
  "203474":"2003474", "203475":"2003475", "203476":"2003476", "203479":"2003479",
  "203480":"2003480", "203482":"2003482", "203486":"2003486", "203506":"2003506",
  "203510":"2003510", "203676":"2003676", "203833":"2003833", "203850":"2003850",
  "203900":"2003900", "203910":"2003910", "203980":"2003980", "204156":"2004156",
  "204157":"2004157", "204163":"2004163", "204174":"2004174", "204263":"2004263",
  "204265":"2004265", "204293":"2004293", "204335":"2004335", "204356":"2004356",
  "204357":"2004357", "204358":"2004358", "204359":"2004359", "204394":"2004394",
  "204397":"2004397", "204572":"2004572", "204665":"2004665", "204765":"2004765",
  "204787":"2004787", "204822":"2004822", "204825":"2004825", "204826":"2004826",
  "204827":"2004827", "204828":"2004828", "204920":"2004920", "204929":"2004929",
  "204963":"2004963", "205018":"2005018", "205019":"2005019", "205043":"2005043",
  "205049":"2005049", "205097":"2005097", "205101":"2005101", "205103":"2005103",
  "205104":"2005104", "205105":"2005105", "205106":"2005106", "205181":"2005181",
  "205203":"2005203", "205204":"2005204", "205205":"2005205", "205206":"2005206",
  "205207":"2005207", "205208":"2005208", "205209":"2005209", "205210":"2005210",
  "205264":"2005264", "205327":"2005327", "205345":"2005345", "205388":"2005388",
  "205457":"2005457", "205461":"2005461", "205499":"2005499", "205502":"2005502",
  "205503":"2005503", "205507":"2005507", "205512":"2005512", "205622":"2005622",
  "205624":"2005624", "205630":"2005630", "205649":"2005649", "205650":"2005650",
  "205651":"2005651", "205652":"2005652", "205672":"2005672", "205689":"2005689",
  "205747":"2005747", "205757":"2005757", "205803":"2005803", "205805":"2005805",
  "205808":"2005808", "205809":"2005809", "205811":"2005811", "205812":"2005812",
  "205813":"2005813", "205816":"2005816", "205905":"2005905", "206345":"2006345",
  "206373":"2006373", "206374":"2006374", "206375":"2006375", "206377":"2006377",
  "206401":"2006401", "206417":"2006417", "206526":"2006526", "206529":"2006529",
  "206583":"2006583", "206584":"2006584", "20668":"2000668", "206725":"2006725",
  "206957":"2006957", "20698":"2000698", "207014":"2007014", "207021":"2007021",
  "207022":"2007022", "207151":"2007151", "207224":"2007224", "207225":"2007225",
  "207230":"2007230", "207231":"2007231", "207232":"2007232", "207283":"2007283",
  "207284":"2007284", "207285":"2007285", "207286":"2007286", "207493":"2007493",
  "207495":"2007495", "207496":"2007496", "207497":"2007497", "207498":"2007498",
  "207499":"2007499", "207500":"2007500", "207501":"2007501", "207502":"2007502",
  "207503":"2007503", "207504":"2007504", "207505":"2007505", "207590":"2007590",
  "207600":"2007600", "207723":"2007723", "207786":"2007786", "207787":"2007787",
  "207788":"2007788", "207789":"2007789", "207790":"2007790", "207791":"2007791",
  "207865":"2007865", "208073":"2008073", "208074":"2008074", "208117":"2008117",
  "208147":"2008147", "208182":"2008182", "208254":"2008254", "208255":"2008255",
  "208256":"2008256", "208257":"2008257", "208260":"2008260", "208269":"2008269",
  "208282":"2008282", "208283":"2008283", "208284":"2008284", "208349":"2008349",
  "208383":"2008383", "208392":"2008392", "208452":"2008452", "208459":"2008459",
  "208460":"2008460", "208461":"2008461", "208599":"2008599", "208833":"2008833",
  "209039":"2009039", "209040":"2009040", "209082":"2009082", "209086":"2009086",
  "209088":"2009088", "209089":"2009089", "209149":"2009149", "209196":"2009196",
  "209198":"2009198", "209207":"2009207", "209210":"2009210", "209231":"2009231",
  "209234":"2009234", "209254":"2009254", "209256":"2009256", "209257":"2009257",
  "209288":"2009288", "209369":"2009369", "209371":"2009371", "209373":"2009373",
  "209374":"2009374", "209393":"2009393", "209394":"2009394", "209514":"2009514",
  "209515":"2009515", "209558":"2009558", "209805":"2009805", "209806":"2009806",
  "209902":"2009902", "209985":"2009985"}


# rules object
rules = []

def getCustomerName(apikey,endpoint,customerid):
    url = endpoint + '/customer/' + customerid
    headers = {'Accept': 'application/json', 'Fastly-Key': apikey}
    r = requests.get(url, headers=headers)
    response = json.loads(r.text)

    return response['name']

def getConfigSet(apikey,endpoint,wafid,version):
    payload = {'include': 'configuration_set'}
    url = endpoint + '/wafs/' + wafid
    headers = {'Accept': 'application/json', 'Fastly-Key': apikey}
    r = requests.get(url, headers=headers, params=payload)
    response = json.loads(r.text)

    configsetid = response['data']['relationships']['configuration_set']['data']['id']

    url = endpoint + '/wafs/configuration_sets'
    headers = {'Accept': 'application/json', 'Fastly-Key': apikey}
    r = requests.get(url, headers=headers)
    response = json.loads(r.text)

    for configset in response['data']:
        if configsetid == configset['id']:
            configsetname = configset['attributes']['name']
            break

    return configsetid,configsetname




def getServiceInfo(apikey,endpoint,serviceid):
    url = endpoint + '/service/' + serviceid + '/details'
    headers = {'Accept': 'application/json', 'Fastly-Key': apikey}
    r = requests.get(url, headers=headers)
    response = json.loads(r.text)

    if 'name' in response:
        version = response['active_version']['number']
        wafid = response['active_version']['wafs'][0]['id']
        customername = getCustomerName(apikey,endpoint,response['customer_id'])
        configsetid, configsetname = getConfigSet(apikey,endpoint,wafid,version)
        print """Operating on Customer: {2}
        - Service: {0} - {1}
        - ConfigSet: {3} - {4}""".format(serviceid,response['name'],customername,configsetid,configsetname)
    else:
        print "Error with service: ", serviceid
        print response
        sys.exit()
    #print json.dumps(response, indent=4, sort_keys=True)
    return version,wafid


def getRules(apikey,endpoint,serviceid,version,wafid):
    url = endpoint + '/service/' + serviceid + '/wafs/' + wafid + '/rule_statuses'
    headers = {'Accept': 'application/vnd.api+json', 'Fastly-Key': apikey}
    r = requests.get(url, headers=headers)
    response = json.loads(r.text)

    # store first page of rules and then iterate through the rest
    rules.append(response['data'])
    current_page = response['meta']['current_page']
    total_pages = response['meta']['total_pages']
    Printer("Grabbing {0} out of {1} pages of rules".format(current_page,total_pages))

    for current_page in range(2, total_pages+1, 1):
        payload = {'page[number]': current_page}
        r = requests.get(url, headers=headers, params=payload)
        response = json.loads(r.text)
        rules.append(response['data'])
        Printer("Grabbing {0} out of {1} pages of rules".format(current_page,total_pages))
        time.sleep(1)
    return rules

    #print json.dumps(response, indent=4, sort_keys=True)

def checkingRules(rules):
    renumbered = []
    print "comparing with updated rules"

    for rule in rules:
        for r in rule:
            if r['attributes']['status'] != 'disabled':
#                print json.dumps(r, indent=4, sort_keys=True)
                current = r['attributes']['modsec_rule_id']
                if current in tw_mapping:
                    renumbered.append(tw_mapping[current])
                    print "we have a match for: {0} and the new rule is {1}".format(current,tw_mapping[current])

    return renumbered


if __name__ == '__main__':

    toolname = os.path.basename(__file__)

    # grab arguments
    parser = argparse.ArgumentParser(description="update TW rule that has been renumbered",
            epilog="""Examples:\n\n
            ** print rules that have been renumbered **\n
            python {0} -k $FASTLY_TOKEN -s <serviceid>\n
            ** print renumbered rules and update config set **\n
            python {0} -k $FASTLY_TOKEN -s <serviceid> -c <new configset ID>\n
            ** print renumbered rules, update config set and re-add trustwave
            rules to new configset with same rule status **\n
            python {0} -k $FASTLY_TOKEN -s <serviceid> -c <new configset ID> -u\n
            """.format(toolname), formatter_class=RawTextHelpFormatter)
    parser.add_argument("-k", "--apikey", required=True, help="API for Fastly")
    parser.add_argument("-e", "--endpoint", required=False, help="API endpoint to be use for fastly", default="https://api.fastly.com")
    parser.add_argument("-s", "--serviceid", required=True, help="Service ID to check")
    parser.add_argument("-c", "--configsetid", required=False, help="configuration set ID to update the WAF to")
    parser.add_argument("-u", "--update", required=False, help="updates the Trustwave rules that have been renumbered on the new config set with the same status, if not set it just prints them")

    # parse them
    args = parser.parse_args()

    apikey = args.apikey
    endpoint = args.endpoint
    serviceid = args.serviceid

    version,wafid = getServiceInfo(apikey,endpoint,serviceid)
    rules = getRules(apikey,endpoint,serviceid,version,wafid)
    renumbered = checkingRules(rules)



    #getWAFID(args.apikey,args.endpoint,args.serviceid,version)
    print "Completed.."