bluethundr12

Untitled

Jan 13th, 2021
1,030
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #!/usr/bin/env python3
  2. #-*- coding: utf-8 -*-
  3. # Import module
  4. import io
  5. import os
  6. import time
  7. import pandas
  8. import argparse
  9. import csv
  10. from pandas import ExcelWriter
  11. from pymongo import MongoClient, errors
  12. from datetime import datetime
  13. from pathlib import Path
  14. from colorama import init, Fore
  15. init()
  16.  
  17.  
  18. ### Generic Functions
  19.  
  20. def welcomebanner():
  21.     # Print the welcome banner
  22.     print(Fore.CYAN)
  23.     message = "*             EC2 MongoDB                     *"
  24.     banner(message, "*")
  25.     print(Fore.RESET)
  26.  
  27. def endbanner():
  28.     print(Fore.CYAN)
  29.     message = "* EC2 MongoDB Operations Are Complete *"
  30.     banner(message, "*")
  31.     print(Fore.RESET)
  32.  
  33. def banner(message, border="-"):
  34.     line = border * len(message)
  35.     print(line)
  36.     print(message)
  37.     print(line)
  38.  
  39. def is_digit(check_input):
  40.     if check_input.isdigit():
  41.         return True
  42.     return False
  43.  
  44. def initialize(interactive, aws_account):
  45.     # Set the date
  46.     today = datetime.today()
  47.     today = today.strftime("%m-%d-%Y")
  48.     # Set the input file
  49.     aws_env_list = os.path.join('..', '..', 'source_files', 'aws_accounts_list', 'aws_accounts_list.csv')
  50.     return today, aws_env_list
  51.  
  52. def connect_db():
  53.     try:
  54.         myclient = MongoClient(
  55.                 host = "mongodb://localhost:27017/",
  56.                 serverSelectionTimeoutMS = 3000 # 3 second timeout
  57.             )
  58.     except errors.ServerSelectionTimeoutError as e:
  59.         # set the client instance to 'None' if exception
  60.         myclient = None
  61.         # catch pymongo.errors.ServerSelectionTimeoutError
  62.         print ("pymongo ERROR:", e)
  63.     return myclient
  64.  
  65. def set_db(instance_col=None):
  66.     if __name__ == "__main__":
  67.         message = "* Select a MongoDB Database *"
  68.         print(Fore.CYAN)
  69.         banner(message, "*")
  70.         print(Fore.RESET)
  71.     myclient = connect_db()
  72.     today = datetime.today()
  73.     today = today.strftime("%m%d%Y")
  74.     if __name__ == "__main__":
  75.         print(Fore.CYAN + "Available MongoDB Databases:")
  76.         if myclient != None:
  77.             # the list_database_names() method returns a list of strings
  78.             database_names = myclient.list_database_names()
  79.             counter = 1
  80.             for db in database_names:
  81.                 message = str(counter) + ". " + db
  82.                 print(message)
  83.                 counter = counter + 1
  84.         print ("There are", len(database_names), "databases.\n")
  85.         print(f"Please select a database. Enter a number 1 through {len(database_names)}.")
  86.         choice = input("Enter a number: ")
  87.         if is_digit(choice) == True:
  88.             if int(choice) > counter:
  89.                 print("Wrong selection.")
  90.                 set_db()
  91.             choice = int(choice)
  92.             choice = choice - 1
  93.             mydb = myclient[database_names[choice]]
  94.             mydb_name = database_names[choice]
  95.             instance_col = "ec2_list_" + today
  96.             instance_col = mydb[instance_col]
  97.             print(f"You've selected: {database_names[choice]}\n")
  98.         else:
  99.             print("Must enter a digit. Try again.\n")
  100.     else:
  101.         mydb = myclient["aws_inventories"]
  102.         mydb_name = "aws_inventories"
  103.         instance_col = "ec2_list_" + today
  104.         instance_col = mydb[instance_col]
  105.     return mydb, mydb_name, instance_col
  106.  
  107. def read_account_info(aws_env_list):
  108.     account_names = []
  109.     account_numbers = []
  110.     with open(aws_env_list) as csv_file:
  111.         csv_reader = csv.reader(csv_file, delimiter=',')
  112.         next(csv_reader)
  113.         for row in csv_reader:
  114.                 account_name = str(row[0])
  115.                 account_number = str(row[1])
  116.                 account_names.append(account_name)
  117.                 account_numbers.append(account_number)
  118.     return account_names, account_numbers
  119.  
  120. def select_account(options, aws_env_list):
  121.     ## Select the account
  122.     if options.account_name:
  123.         aws_account = options.account_name
  124.     else:
  125.         print(Fore.YELLOW)
  126.         aws_account = input("Enter the name of the AWS account you'll be working in: ")
  127.         print(Fore.RESET)
  128.     aws_account_number = find_account_number(aws_account, aws_env_list)
  129.     return aws_account, aws_account_number
  130.  
  131. def find_account_number(aws_account,aws_env_list):
  132.     account_names, account_numbers = read_account_info(aws_env_list)
  133.     for (my_aws_account, my_aws_account_number) in zip(account_names, account_numbers):
  134.         if my_aws_account == aws_account:
  135.             aws_account_number = my_aws_account_number
  136.     if aws_account ==  "all":
  137.         aws_account_number = '1234567891011'
  138.     return aws_account_number
  139.  
  140. def arguments():
  141.     parser = argparse.ArgumentParser(description="This is a program that provides a text interface to MongoDB.")
  142.  
  143.     parser.add_argument(
  144.     "-n",
  145.     "--account_name",
  146.     type = str,
  147.     default = None,
  148.     nargs = "?",
  149.     help = "Name of the AWS account you'll be working in")
  150.  
  151.     options = parser.parse_args()
  152.     return options
  153.  
  154. def create_directories():
  155.  
  156.     ## Set source and output file directories
  157.     source_files_path = Path('..', '..', 'source_files', 'aws_accounts_list')
  158.     output_files_path = Path('..', '..', 'output_files', 'aws_instance_list')
  159.  
  160.     os.makedirs(source_files_path, exist_ok=True)
  161.     os.makedirs(output_files_path, exist_ok=True)
  162.  
  163.     # Create output subdirectories
  164.     folders = ['csv','excel','html', 'json']
  165.     for folder in folders:
  166.         full_path = os.path.join(output_files_path,folder)
  167.         os.makedirs(full_path, exist_ok=True)
  168.  
  169. def set_test_dict():
  170.     mydict = { "AWS Account": "company-lab", "Account Number": "12345678910", "Name": "bastion001",
  171. "Instance ID": "i-07aaef3b7167d592a", "AMI ID": "ami-07fd81f1ecf6cf387", "Volumes": "vol-09d6d898db4af132a",
  172. "Private IP": "10.238.3.165", "Public IP": "xx.xx.xx.xx", "Private DNS": "ip-10-238-3-165.ec2.internal",
  173. "Availability Zone": "us-east-1a", "VPC ID": "vpc-00de11103235ec567", "Type": "t3.small", "Key Pair Name": "ccmi-vzn-int01", "Instance State": "running", "Launch Date": "September 10 2019"}
  174.     return mydict
  175.  
  176. ### Main functions
  177.  
  178. # Choice 1. Create new MongoDB Database
  179. def create_mongodb(mydict):
  180.     myclient = connect_db()
  181.     message = f"* Create new MongoDB *"
  182.     banner(message, border="*")
  183.     print("\n")
  184.     newdb = input("Enter the name of a new mongo database: ")
  185.     dblist = myclient.list_database_names()
  186.     if newdb in dblist:
  187.         print("The database exists.")
  188.         main()
  189.     else:
  190.         try:
  191.             mydb = myclient[newdb]
  192.             mycol = mydb["testColumn"]
  193.             mycol.insert_one(mydict)
  194.             message = f"Succeeded in creating: {newdb}"
  195.             banner(message)
  196.         except Exception as e:
  197.             print(f"MongoDB Database creation failed with: {e}")
  198.  
  199. # Choice 2. Drop MongoDB Database
  200. def drop_mongodb():
  201.     message = "* Drop MongoDB *"
  202.     banner(message, "*")
  203.     myclient = connect_db()
  204.     today = datetime.today()
  205.     today = today.strftime("%m%d%Y")
  206.     if myclient != None:
  207.         # the list_database_names() method returns a list of strings
  208.         database_names = myclient.list_database_names()
  209.         counter = 1
  210.         for db in database_names:
  211.             message = str(counter) + ". " + db
  212.             print(message)
  213.             counter = counter + 1
  214.     print ("There are", len(database_names), "databases.\n")
  215.     db_names_before_drop = myclient.list_database_names()
  216.     print ("db count BEFORE drop:", len(db_names_before_drop))
  217.     print(f"Please select a database. Enter a number 1 through {len(database_names)}.")
  218.     choice = input("Enter a number: ")
  219.     if is_digit(choice) == True:
  220.         if int(choice) > counter:
  221.             print("Wrong selection.")
  222.             set_db()
  223.         choice = int(choice)
  224.         choice = choice - 1
  225.         dropdb = myclient[database_names[choice]]
  226.         instance_col = "ec2_list_" + today
  227.         instance_col = dropdb[instance_col]
  228.         print(f"You've selected: {database_names[choice]}\n")
  229.     else:
  230.         print("Must enter a digit. Try again.\n")
  231.     # check if a collection exists
  232.     col_exists = instance_col in dropdb.list_collection_names()
  233.     print ("Some Collection exists:", col_exists) # will print True or False
  234.     # call MongoDB client object"s drop_database() method to delete a db
  235.     myclient.drop_database(dropdb) # pass db name as string
  236.     time.sleep(5)
  237.     # get all of the database names
  238.     db_names_after_drop = myclient.list_database_names()
  239.     print ("db count AFTER drop:", len(db_names_before_drop))
  240.     diff = len(db_names_before_drop) - len(db_names_after_drop)
  241.     print ("difference:", diff)
  242.  
  243. # Choice 3. Insert into MongoDB
  244. def insert_doc(mydict):
  245.     _, _, instance_col = set_db()
  246.     mydict["_id"] = ObjectId()
  247.     instance_doc = instance_col.insert_one(mydict)
  248.     if __name__ == "__main__":
  249.         message = "* MongoDB Insert Document *"
  250.         banner(message, "*")
  251.         message = f"MongoDB record inserted: {instance_doc.inserted_id}"
  252.         banner(message)
  253.     return instance_doc
  254.  
  255. # Choice 4. Clear the DB
  256. def clear_db():
  257.     _, _, instance_col = set_db()
  258.     message = f"* Clear the DB *"
  259.     banner(message, border="*")
  260.     print(f"This command empties the database.\n")
  261.     try:
  262.         x = instance_col.delete_many({})
  263.     except Exception as e:
  264.         print(f"An error has occurred: {e}")
  265.     print(x.deleted_count, "documents deleted.")
  266.  
  267.  
  268. # Choice 5. Remove Accounts from the DB
  269. def delete_from_collection(aws_account_number):
  270.     _, _, instance_col = set_db()
  271.     if __name__ == "__main__":
  272.         message = f"* Clear old entries *"
  273.         banner(message, border="*")
  274.         print(f"This command clears old entries the database.\n")
  275.         aws_account_number = input("Enter an AWS account number: ")
  276.     try:
  277.         #instance_col.remove({"Account Number": aws_account_number});
  278.         instance_col.delete_many({"Account Number": aws_account_number})
  279.     except Exception as e:
  280.         print(f"An error has occurred: {e}")
  281.  
  282. # Choice 6. Print the DB
  283. def mongo_select_all():
  284.     _, mydb_name, instance_col = set_db()
  285.     instance_list = list(instance_col.find())
  286.     if __name__ == "__main__":
  287.         message = f"* Print DB Documents in {mydb_name} *"
  288.         banner(message, border="*")
  289.         print("\n")
  290.         if not instance_list:
  291.             message = f"The database: {mydb_name} has no entries."
  292.             banner(message)
  293.         else:
  294.             message = f"The databse: {mydb_name} has {len(instance_list)} entries."
  295.             for data in instance_list:
  296.                 print(data)
  297.     print("\n")
  298.     return instance_list
  299.  
  300. # Choice 7. Print DB Names
  301. def print_db_names():
  302.     myclient = connect_db()
  303.     message = f"* Print DB Names *"
  304.     banner(message, border="*")
  305.     print("The database names are:")
  306.     if myclient != None:
  307.         # the list_database_names() method returns a list of strings
  308.         database_names = myclient.list_database_names()
  309.         counter = 1
  310.         for db in database_names:
  311.             message = str(counter) + ". " + db
  312.             print(message)
  313.             counter = counter + 1
  314.         print ("There are", len(database_names), "databases.")
  315.  
  316. # Choice 8. Print Collections
  317. def print_collections():
  318.     myclient = connect_db()
  319.     message = f"* Print DB Collections *"
  320.     banner(message, border="*")
  321.     print(f"This command prints the database collection names.\n")
  322.     if myclient != None:
  323.         # the list_database_names() method returns a list of strings
  324.         database_names = myclient.list_database_names()
  325.         print ("There are", len(database_names), "databases.")
  326.         for db_num, db in enumerate(database_names):
  327.             print ("\nGetting collections for database:", db, "--", db_num)
  328.             collection_names = myclient[db].list_collection_names()
  329.             print ("The MongoDB database returned", len(collection_names), "collections.")
  330.             # iterate over the list of collection names
  331.             for col_num, col in enumerate(collection_names):
  332.                 print (col, "--", col_num)
  333.  
  334. # Choice 9. Export MongoDB to file
  335. def mongo_export_to_file(interactive, aws_account, aws_account_number,instance_col=None,date=None):
  336.     create_directories()
  337.     if date == None:
  338.         today = datetime.today()
  339.         today = today.strftime("%m-%d-%Y")
  340.         date = today
  341.     #today = datetime.today()
  342.     #date = date.strftime("%m-%d-%Y")
  343.     if not instance_col:
  344.         _, _, instance_col = set_db()
  345.     # make an API call to the MongoDB server
  346.     if interactive == 0:
  347.         mongo_docs = instance_col.find({})
  348.     else:
  349.         mongo_docs = instance_col.find({"Account Number": aws_account_number})
  350.     # Convert the mongo docs to a DataFrame
  351.     docs = pandas.DataFrame(mongo_docs)
  352.     # Discard the Mongo ID for the documents
  353.     docs.pop("_id")
  354.     if __name__ == "__main__":
  355.         print("Choose a file format")
  356.         print("1. CSV")
  357.         print("2. JSON")
  358.         print("3. HTML")
  359.         print("4. Excel")
  360.         choice = input("Enter a number 1-4: ")
  361.         choice = int(choice)
  362.     else:
  363.         choice = 1
  364.     if choice == 1:
  365.         if __name__ == "__main__":
  366.             # export MongoDB documents to CSV
  367.             csv_export = docs.to_csv(sep=",") # CSV delimited by commas
  368.             print ("\nCSV data:", csv_export)
  369.         # Set the CSV output directory
  370.         output_dir = os.path.join("..", "..", "output_files", "aws_instance_list", "csv", "")
  371.         if interactive == 1:
  372.             output_file = os.path.join(output_dir, "aws-instance-list-" + aws_account + "-" + date +".csv")
  373.         else:
  374.             output_file = os.path.join(output_dir, "aws-instance-master-list-" + date +".csv")
  375.  
  376.         # export MongoDB documents to a CSV file, leaving out the row "labels" (row numbers)
  377.         docs.to_csv(output_file, ",", index=False) # CSV delimited by commas
  378.     elif choice == 2:
  379.         if __name__ == "__main__":
  380.             json_export = docs.to_json() # return JSON data
  381.             print ("\nJSON data:", json_export)
  382.         # Set the JSON output directory
  383.         output_dir = os.path.join("..", "..", "output_files", "aws_instance_list", "json", "")
  384.         if interactive == 1:
  385.             output_file = os.path.join(output_dir, "aws-instance-list-" + aws_account + "-" + date +".json")
  386.         else:
  387.             output_file = os.path.join(output_dir, "aws-instance-master-list-" + date +".json")
  388.         # export MongoDB documents to a CSV file, leaving out the row "labels" (row numbers)
  389.         docs.to_json(output_file)
  390.     elif choice == 3:
  391.         html_str = io.StringIO()
  392.         # export as HTML
  393.         docs.to_html(
  394.         buf=html_str,
  395.         classes="table table-striped"
  396.         )
  397.         if __name__ == "__main__":
  398.             # print out the HTML table
  399.             print (html_str.getvalue())
  400.         # Set the HTML output directory
  401.         output_dir = os.path.join("..", "..", "output_files", "aws_instance_list", "html", "")
  402.         if interactive == 1:
  403.             output_file = os.path.join(output_dir, "aws-instance-list-" + aws_account + "-" + date +".html")
  404.         else:
  405.             output_file = os.path.join(output_dir, "aws-instance-master-list-" + date + ".html")
  406.         # save the MongoDB documents as an HTML table
  407.         docs.to_html(output_file)
  408.     elif choice == 4:
  409.         # Set the Excel output directory
  410.         output_dir = os.path.join("..", "..", "output_files", "aws_instance_list", "excel", "")
  411.         time.sleep(5)
  412.         if interactive == 1:
  413.             output_file = os.path.join(output_dir, "aws-instance-list-" + aws_account + "-" + date + ".xlsx")
  414.         else:
  415.             output_file = os.path.join(output_dir, "aws-instance-master-list-" + date + ".xlsx")
  416.         # export MongoDB documents to a Excel file, leaving out the row "labels" (row numbers)
  417.         writer = ExcelWriter(output_file)
  418.         docs.to_excel(writer,"EC2 List",index=False)
  419.         writer.save()
  420.     if __name__ == "__main__":
  421.         exit = input("Exit program (y/n): ")
  422.         if exit.lower() == "y" or exit.lower() == "yes":
  423.             exit_program()
  424.         else:
  425.             main()
  426.  
  427. # Choice 10. Print Reports
  428. def print_reports(interactive,aws_account,aws_account_number):
  429.     set_db(instance_col=None)
  430.     inputDate = input("Enter the date in format 'dd/mm/yyyy': ")
  431.     day,month,year = inputDate.split('/')
  432.     isValidDate = True
  433.     try:
  434.         datetime(int(year),int(month),int(day))
  435.     except ValueError :
  436.         isValidDate = False
  437.         print_reports(interactive,aws_account,aws_account_number)
  438.  
  439.     if(isValidDate) :
  440.         print(f"Input date is valid: {inputDate}")
  441.         format= "%m%d%Y"
  442.         inputDate = datetime.strptime(inputDate,"%m/%d/%Y")
  443.         inputDate = inputDate.strftime(format)
  444.     else:
  445.         print(f"Input date is not valid: {inputDate}")
  446.         print_reports(interactive,aws_account,aws_account_number)
  447.     myclient = connect_db()
  448.     mydb = myclient["aws_inventories"]
  449.     instance_col = "ec2_list_" + inputDate
  450.     instance_col = mydb[instance_col]
  451.     mongo_export_to_file(interactive, aws_account, aws_account_number,instance_col,date=inputDate)
  452.  
  453. # Choice 11. Exit ec2 Mongo
  454. def exit_program():
  455.     endbanner()
  456.     exit()
  457.  
  458. def menu():
  459.     message = "Main Menu"
  460.     banner(message)
  461.     print(Fore.CYAN + "Your available actions: ")
  462.     print("1. Create new MongoDB Database")
  463.     print("2. Drop MongoDB Database")
  464.     print("3. Do a test insert to the DB")
  465.     print("4. Clear the DB")
  466.     print("5. Remove accounts from the DB.")
  467.     print("6. Print the DB")
  468.     print("7. Print DB Names")
  469.     print("8. Print collections")
  470.     print("9. Export MongoDB to file")
  471.     print("10. Print Reports")
  472.     print("11. Exit ec2 mongo")
  473.     print("\n")
  474.  
  475. def main():
  476.     options = arguments()
  477.     welcomebanner()
  478.     mydict = set_test_dict()
  479.     if __name__ == "__main__":
  480.         print(Fore.YELLOW)
  481.         aws_accounts_answer = input("Work in one or all accounts: ")
  482.         print(Fore.RESET)
  483.         # Set interacive variable to indicate one or many accounts
  484.         if aws_accounts_answer.lower() == "one" or aws_accounts_answer.lower() == "1":
  485.             interactive = 1
  486.         else:
  487.             interactive = 0
  488.         aws_account = ''
  489.         _, aws_env_list = initialize(interactive, aws_account)
  490.         menu()
  491.         option = input("Enter the option: ")
  492.         option = int(option)
  493.         # 1. Create a MongoDB database
  494.         if option == 1:
  495.             create_mongodb(mydict)
  496.             main()
  497.         # 2. Drop a MongoDB Database
  498.         elif option == 2:
  499.             drop_mongodb()
  500.             main()
  501.         # 3. Do a test insert to the DB
  502.         elif option  == 3:
  503.             insert_doc(mydict)
  504.             main()
  505.         # 4. Clear the DB"
  506.         elif option == 4:
  507.             clear_db()
  508.             main()
  509.         # 5. Remove accounts from the DB.
  510.         elif option == 5:
  511.             _, aws_account_number = select_account(options, aws_env_list)
  512.             delete_from_collection(aws_account_number)
  513.             main()
  514.         # 6. Print the DB
  515.         elif option == 6:
  516.             mongo_select_all()
  517.             main()
  518.         # 7. Print DB Names
  519.         elif option == 7:
  520.             print_db_names()
  521.             main()
  522.         # 8. Print collections
  523.         elif option == 8:
  524.             print_collections()
  525.             main()
  526.         # 9. Export MongoDB to file
  527.         elif option == 9:
  528.             if aws_accounts_answer == "all":
  529.                 aws_account = "all"
  530.                 aws_account_number = "123456789101"
  531.             else:
  532.                 aws_account, aws_account_number = select_account(options, aws_env_list)
  533.             mongo_export_to_file(interactive, aws_account, aws_account_number)
  534.             main()
  535.         # 10 Print Reports
  536.         elif option == 10:
  537.             if aws_accounts_answer == "all":
  538.                 aws_account = "all"
  539.                 aws_account_number = "123456789101"
  540.             else:
  541.                 aws_account, aws_account_number = select_account(options, aws_env_list)
  542.             print_reports(interactive,aws_account,aws_account_number)
  543.         # 11. Exit ec2 mongo
  544.         elif option == 11:
  545.             exit_program()
  546.         # Invalid Input
  547.         else:
  548.             message = "That is not a valid option."
  549.             banner(message)
  550.             main()
  551.  
  552. if __name__ == "__main__":
  553.     main()
RAW Paste Data