OS-47

main.py

@@ -1,10 +1,40 @@
 import argparse
 from art import text2art
-
+import random
+import boto3
+import os
+import glob
 from src.logger import setup_logger
 from src.snapper import Snapper
 from src.scanner import Scanner
 
+
+def getting_all_pem_file_names():
+    """
+    :return: .pem file names from the red-detector directory.
+    """
+    file_path = os.path.realpath(__file__)  # getting the script's path
+    file_path = file_path.split("red-detector")
+    files_path = file_path[0] + "red-detector"  # (the pem files arent in the same directory as the script.)
+
+    lst = (glob.glob(files_path+"/*.pem"))
+    index = 0
+    for i in lst:
+        lst[index] = lst[index].replace(files_path+"/", "").replace(".pem","")
+        index += 1
+    return lst
+
+
+def used_key_pairs():
+    keypairs = []  # list of used keyPair names
+    ec2 = boto3.client('ec2')
+    response = ec2.describe_key_pairs()
+
+    for i in response["KeyPairs"]:
+        keypairs.append(i["KeyName"])
+    return keypairs
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument('--region', action='store', dest='region', type=str,
@@ -37,11 +67,20 @@
 
     volume_id, selected_az, snapshot_id = snapper.snapshot2volume(volume_id=source_volume_id)
 
-    scanner = Scanner(logger=logger, region=snapper.region)
     if cmd_args.keypair:
-        scanner.keypair_name = cmd_args.keypair
+        scanner = Scanner(logger=logger, region=snapper.region, key_pair_name=cmd_args.keypair)
     else:
-        scanner.keypair_name = scanner.create_keypair(key_name='red_detector_key')
+        used_key_pairs_list_from_aws = used_key_pairs()
+        used_key_pairs_list_locally = getting_all_pem_file_names()
+        num = 0
+        key_name = "red_detector_key{number}".format(number=str(num))
+        while key_name in used_key_pairs_list_from_aws or key_name in used_key_pairs_list_locally:
+            num += 1
+            key_name = "red_detector_key{number}".format(number=str(num))
+
+        scanner = Scanner(logger=logger, region=snapper.region, key_pair_name=key_name)
+        scanner.keypair_name = scanner.create_keypair(key_name=key_name)
+
     ec2_instance_id, ec2_instance_public_ip, report_service_port = scanner.create_ec2(selected_az=selected_az)
     scanner.attach_volume_to_ec2(ec2_instance_id=ec2_instance_id, volume_id=volume_id)
     scanner.scan_and_report(ec2_instance_public_ip=ec2_instance_public_ip,

src/scanner.py

@@ -3,6 +3,7 @@
 import time
 
 import boto3
+import subprocess
 import paramiko
 import requests
 from botocore.exceptions import ClientError, WaiterError
@@ -12,14 +13,15 @@
 
 
 class Scanner:
-    def __init__(self, logger, region):
+    def __init__(self, logger, region, key_pair_name):
         self.logger = logger
         self.region = region
+        self.key_pair_name = key_pair_name
         self.client = boto3.client('ec2', region_name=region)
         self.ec2 = boto3.resource('ec2', region_name=region)
         self.keypair_name = None
 
-    def create_keypair(self, key_name='red_detector_key'):
+    def create_keypair(self, key_name):
         try:
             new_keypair = self.ec2.create_key_pair(KeyName=key_name)
         except ClientError as err:
@@ -30,9 +32,10 @@ def create_keypair(self, key_name='red_detector_key'):
                     return key_name
             self.logger.error(f"create key pair: {err}")
             exit(99)
-        self.logger.info(f'creating key pair: "red_detector_key"')
-        with open('red_detector_key.pem', 'w') as f:
+        self.logger.info('creating key pair: {red_detector_key}'.format(red_detector_key=self.key_pair_name))
+        with open(self.key_pair_name+'.pem', 'w') as f:
             f.write(new_keypair.key_material)
+            output = subprocess.getoutput("chmod 400 "+self.key_pair_name+'.pem')
         return key_name
 
     @staticmethod
@@ -137,7 +140,7 @@ def create_ec2(self, selected_az):
                 MinCount=1,
                 MaxCount=1,
                 InstanceType='t2.large',
-                KeyName=self.keypair_name,
+                KeyName=self.key_pair_name,
                 UserData=user_data,
                 SecurityGroupIds=[
                     security_group_id,
@@ -208,7 +211,7 @@ def attach_volume_to_ec2(self, ec2_instance_id, volume_id):
     def scan_and_report(self, ec2_instance_public_ip, report_service_port, ec2_instance_id, snapshot_id):
         ssh = paramiko.SSHClient()
         ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-        privet_key = paramiko.RSAKey.from_private_key_file("red_detector_key.pem")
+        privet_key = paramiko.RSAKey.from_private_key_file(self.key_pair_name+".pem")
         connect = 0
         while not connect:
             try: