diff --git a/main.py b/main.py index 1c3070e..2b93323 100644 --- a/main.py +++ b/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, diff --git a/src/scanner.py b/src/scanner.py index 108831e..178349d 100644 --- a/src/scanner.py +++ b/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: