bluethundr12

Untitled

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

Adblocker detected! Please consider disabling it...

We've detected AdBlock Plus or some other adblocking software preventing Pastebin.com from fully loading.

We don't have any obnoxious sound, or popup ads, we actively block these annoying types of ads!

Please add Pastebin.com to your ad blocker whitelist or disable your adblocking software.

×