0
votes

I am currently writing a script for my library, and i ended up getting stuck on how to design argparse that will have a lot of options, and sub arguments.

Currently i am designing the search function, which has the following options, some required and some not:

  • search_session_id - Required
  • user_session_id - Required
  • discover_fields - Optional
  • start_time - Optional
  • end_time - Optional
  • summary_fields - Optional
  • field_summary - Optional
  • local_search - Optional

My problem is then the following:

How can i make the argparse and the if statements, if all optionals needs to work together, but also work if only one of them are defined?

If i need to check every single combination, i would end up with something like this:

#!/usr/bin/env python3

"""Script to generate searches on the ArcSight Logger"""

import arcsightrest
import argparse

parser = argparse.ArgumentParser(description='Script used to send search '
                                             'queries to ArcSight Logger API')
parser.add_argument('-t', '--target',
                    help='IP Address of the Loggger', required=True)
parser.add_argument('-u', '--username',
                    help='Username to access the logger', required=True)
parser.add_argument('-p', '--password',
                    help='Password to access the logger', required=True)
parser.add_argument('-ussl', '--unsecuressl', action='store_true',
                    help='Disable ssl warnings', )
parser.add_argument('-w', '--wait', action='store_true',
                    help='Wait for query to finish', )
parser.add_argument('-q', '--query',
                    help='Query to be used in the search')
parser.add_argument('-st', '--starttime',
                    help='From which time the query should look')
parser.add_argument('-et', '--endtime',
                    help='To which time the query should look')
parser.add_argument('-e', '--event',
                    help='Events based input search id')
parser.add_argument('-s', '--status',
                    help='Status of running search')
args = (parser.parse_args())

"""
Sets the target Logger Server
"""
arcsightrest.ArcsightLogger.TARGET = args.target

"""
Gets login token from the Logger API
"""
arc = arcsightrest.ArcsightLogger(args.username, args.password,
                                  args.unsecuressl)
"""
Checks if query is used, and starts a search
"""
if args.query:
    if args.starttime:
        search_id, response = arc.search(args.query, start_time=args.starttime,
                                         end_time=args.endtime)
    search_id, response = arc.search(args.query)

    if args.starttime and args.discover_fields:
        search_id, response = arc.search(args.query, start_time=args.starttime,
                                         end_time=args.endtime,
                                         discover_fields=args.discover_fields)
    print('The search id is {}'.format(search_id))
    if response:
        print('The search has successfully started')

As you can see, i can continue with no end, to make if statements that have every single combination of optional arguments. There must be an easier way to design this? If i was to just parse it in as kwargs, they would not be sent in the correct format, or i would require the person using the script to write things like end_time=SOMETIME, instead of just --endtime TIME. Now this might seem to be a small price to pay, but if i need to add every function with all their parameters into the script, then this will be come a lot longer and more tedious.

1

1 Answers

2
votes

You could collect all the optional keyword arguments passed to arc.search to a dict and then unpack it when you call the function:

import argparse

parser = argparse.ArgumentParser(description='Script used to send search '
                                             'queries to ArcSight Logger API')
parser.add_argument('-t', '--target',
                    help='IP Address of the Loggger', required=True)
parser.add_argument('-u', '--username',
                    help='Username to access the logger', required=True)
parser.add_argument('-p', '--password',
                    help='Password to access the logger', required=True)
parser.add_argument('-q', '--query',
                    help='Query to be used in the search')
parser.add_argument('-st', '--starttime',
                    help='From which time the query should look')
parser.add_argument('-et', '--endtime',
                    help='To which time the query should look')
args = (parser.parse_args())

# Mock search
def search(query, start_time=None, end_time=None, discover_fields=None):
    return 'Id', ', '.join(str(x) for x in [start_time, end_time, discover_fields])

"""
Checks if query is used, and starts a search
"""
if args.query:
    # {name used in argparse: search parameter name}
    query_args = {
        'starttime': 'start_time',
        'endtime': 'end_time',
        'discover_fields': 'discover_fields'
    }
    d = vars(args)
    real_args = {v: d[k] for k, v in query_args.items() if k in d}
    search_id, response = search(args.query, **real_args)

    print('The search id is {}'.format(search_id))
    print('Response is {}'.format(response))

Output:

>python test.py -t foo -u user -p pass -q
query -st start -et end
The search id is Id
Response is start, end, None

Since some of the argument names used by parser are different than the ones passed to search the names need to remapped. vars is used to create a dict from Namespace object returned by parse_args(). Then dictionary comprehension iterates over the mapped argument names, picks ones which were given user and creates a new dictionary with key names that arc.search understands. Finally **real_args unpacks the dictionary named parameters within the function call.