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
parser = argparse.ArgumentParser(description="Export attack mitigations from a database")
parser.add_argument(
"-host", "--mongo_host", default="localhost", help="URL for mongo database.", required=False
)
parser.add_argument(
"-port",
"--mongo_port",
action="store",
default=27017,
type=int,
help="Port for mongo database. Default 27017",
required=False,
)
parser.add_argument(
"-db",
"--database_name",
action="store",
default="monkeyisland",
help="Database name inside of mongo.",
required=False,
)
parser.add_argument(
"-cn",
"--collection_name",
action="store",
default="attack_mitigations",
help="Which collection are we going to export",
required=False,
)
args = parser.parse_args()
def parse_args():
parser = argparse.ArgumentParser(description="Export attack mitigations from a database")
parser.add_argument(
"-host", "--mongo_host", default="localhost", help="URL for mongo database.", required=False
)
parser.add_argument(
"-port",
"--mongo_port",
action="store",
default=27017,
type=int,
help="Port for mongo database. Default 27017",
required=False,
)
parser.add_argument(
"-db",
"--database_name",
action="store",
default="monkeyisland",
help="Database name inside of mongo.",
required=False,
)
parser.add_argument(
"-cn",
"--collection_name",
action="store",
default="attack_mitigations",
help="Which collection are we going to export",
required=False,
)
return parser.parse_args()
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():
args = parse_args()
mongodb = connect_to_mongo(args.mongo_host, args.mongo_port, args.database_name)
clean_collection(mongodb, args.collection_name)