Deployment: Reorder functions in dump_attack_mitigations.py

This commit is contained in:
Mike Salvatore 2021-09-28 12:02:52 -04:00
parent 38f50641a5
commit 82c8385863
1 changed files with 10 additions and 10 deletions

View File

@ -3,6 +3,13 @@ import argparse
import pymongo import pymongo
def main():
args = parse_args()
mongodb = connect_to_mongo(args.mongo_host, args.mongo_port, args.database_name)
clean_collection(mongodb, args.collection_name)
def parse_args(): def parse_args():
parser = argparse.ArgumentParser(description="Export attack mitigations from a database") parser = argparse.ArgumentParser(description="Export attack mitigations from a database")
parser.add_argument( parser.add_argument(
@ -42,21 +49,14 @@ def connect_to_mongo(mongo_host: str, mongo_port: int, database_name: str) -> py
return database return database
def collection_exists(mongodb: pymongo.MongoClient, collection_name: str) -> bool:
collections = mongodb.list_collection_names()
return collection_name in collections
def clean_collection(mongodb: pymongo.MongoClient, collection_name: str): def clean_collection(mongodb: pymongo.MongoClient, collection_name: str):
if collection_exists(mongodb, collection_name): if collection_exists(mongodb, collection_name):
mongodb.drop_collection(collection_name) mongodb.drop_collection(collection_name)
def main(): def collection_exists(mongodb: pymongo.MongoClient, collection_name: str) -> bool:
args = parse_args() collections = mongodb.list_collection_names()
mongodb = connect_to_mongo(args.mongo_host, args.mongo_port, args.database_name) return collection_name in collections
clean_collection(mongodb, args.collection_name)
if __name__ == "__main__": if __name__ == "__main__":