Deployment: Wrap argument parsing in function

This commit is contained in:
Mike Salvatore 2021-09-28 12:01:16 -04:00
parent 7bcfc6d27a
commit 38f50641a5
1 changed files with 33 additions and 30 deletions

View File

@ -2,36 +2,38 @@ import argparse
import pymongo import pymongo
parser = argparse.ArgumentParser(description="Export attack mitigations from a database")
parser.add_argument( def parse_args():
"-host", "--mongo_host", default="localhost", help="URL for mongo database.", required=False parser = argparse.ArgumentParser(description="Export attack mitigations from a database")
) parser.add_argument(
parser.add_argument( "-host", "--mongo_host", default="localhost", help="URL for mongo database.", required=False
"-port", )
"--mongo_port", parser.add_argument(
action="store", "-port",
default=27017, "--mongo_port",
type=int, action="store",
help="Port for mongo database. Default 27017", default=27017,
required=False, type=int,
) help="Port for mongo database. Default 27017",
parser.add_argument( required=False,
"-db", )
"--database_name", parser.add_argument(
action="store", "-db",
default="monkeyisland", "--database_name",
help="Database name inside of mongo.", action="store",
required=False, default="monkeyisland",
) help="Database name inside of mongo.",
parser.add_argument( required=False,
"-cn", )
"--collection_name", parser.add_argument(
action="store", "-cn",
default="attack_mitigations", "--collection_name",
help="Which collection are we going to export", action="store",
required=False, default="attack_mitigations",
) help="Which collection are we going to export",
args = parser.parse_args() required=False,
)
return parser.parse_args()
def connect_to_mongo(mongo_host: str, mongo_port: int, database_name: str) -> pymongo.MongoClient: def connect_to_mongo(mongo_host: str, mongo_port: int, database_name: str) -> pymongo.MongoClient:
@ -51,6 +53,7 @@ def clean_collection(mongodb: pymongo.MongoClient, collection_name: str):
def main(): def main():
args = parse_args()
mongodb = connect_to_mongo(args.mongo_host, args.mongo_port, args.database_name) mongodb = connect_to_mongo(args.mongo_host, args.mongo_port, args.database_name)
clean_collection(mongodb, args.collection_name) clean_collection(mongodb, args.collection_name)