diff --git a/dojo/management/commands/notify_isoc.py b/dojo/management/commands/notify_isoc.py index 99c4339941..67327fd162 100644 --- a/dojo/management/commands/notify_isoc.py +++ b/dojo/management/commands/notify_isoc.py @@ -20,17 +20,21 @@ class Command(BaseCommand): help = "Details: Spams External Unit\nArgs: Weekly, Monthly, Quarterly" + def add_arguments(self, parser): + parser.add_argument('type') + def handle(self, *args, **options): + type = options['type'] - if not args: + if not options: print "Must specify an argument: Weekly, Monthly, or Quarterly" sys.exit(0) - if args[0] not in ["Weekly", "Monthly", "Quarterly"]: - print("Unexpected frequency: " + str(args[0]) + + if type not in ["Weekly", "Monthly", "Quarterly"]: + print("Unexpected frequency: " + str(type) + "\nMust specify an argument: Weekly, Monthly, or Quarterly.") sys.exit(0) - scSettings = ScanSettings.objects.filter(frequency=args[0]) + scSettings = ScanSettings.objects.filter(frequency=type) scan_start_time = datetime.datetime.today() + datetime.timedelta( hours=12) diff --git a/dojo/management/commands/run_scan.py b/dojo/management/commands/run_scan.py index 196e37e26d..067b60d3e7 100755 --- a/dojo/management/commands/run_scan.py +++ b/dojo/management/commands/run_scan.py @@ -27,7 +27,11 @@ class Command(BaseCommand): help = "Details:\n\tRuns nmap scans\n\nArguments:" +\ "\n\tWeekly\n\tMonthly\n\tQuarterly" + def add_arguments(self, parser): + parser.add_argument('type') + def handle(self, *args, **options): + type = options['type'] # Scan the host and add the results of the scan to the host informaiton def runScan(prod_id, p_dict): @@ -195,12 +199,12 @@ def run(self): Scans are performed on a Weekly, Monthly, or Quarterly bases. The target frequency is specified by the cron job scheduler. """ - if not args: + if not options: print "Must specify an argument: Weekly, Monthly, Quarterly, or ID",\ " of Scan Settings to use." sys.exit(0) - if (args[0] in ["Weekly", "Monthly", "Quarterly"] - or args[0].isdigit()): + if (type in ["Weekly", "Monthly", "Quarterly"] + or type.isdigit()): pass else: print("Unexpected parameter: " + str(args[0])) @@ -208,10 +212,10 @@ def run(self): " or ID of Scan Settings to use." sys.exit(0) - if args[0].isdigit(): - scSettings = ScanSettings.objects.filter(id=args[0]) + if type.isdigit(): + scSettings = ScanSettings.objects.filter(id=type) else: - scSettings = ScanSettings.objects.filter(frequency=args[0]) + scSettings = ScanSettings.objects.filter(frequency=type) if len(scSettings) <= 0: print("No scan settings found with parameter specified.")