Add most infrastrucure for running AWS commands

This commit is contained in:
Itay Mizeretz 2019-02-03 14:18:08 +02:00
parent 1cedfb5c2d
commit 378baa7139
11 changed files with 161 additions and 7 deletions

View File

@ -3,7 +3,7 @@ import urllib2
__author__ = 'itay.mizeretz' __author__ = 'itay.mizeretz'
class AWS(object): class AwsInstance(object):
def __init__(self): def __init__(self):
try: try:
self.instance_id = urllib2.urlopen('http://169.254.169.254/latest/meta-data/instance-id').read() self.instance_id = urllib2.urlopen('http://169.254.169.254/latest/meta-data/instance-id').read()

View File

@ -0,0 +1,41 @@
import boto3
__author__ = 'itay.mizeretz'
class AwsService(object):
"""
Supplies various AWS services
"""
# TODO: consider changing from static to singleton, and generally change design
access_key_id = None
secret_access_key = None
region = None
@staticmethod
def set_auth_params(access_key_id, secret_access_key):
AwsService.access_key_id = access_key_id
AwsService.secret_access_key = secret_access_key
@staticmethod
def set_region(region):
AwsService.region = region
@staticmethod
def get_client(client_type, region=None):
return boto3.client(
client_type,
aws_access_key_id=AwsService.access_key_id,
aws_secret_access_key=AwsService.secret_access_key,
region_name=region if region is not None else AwsService.region)
@staticmethod
def get_session():
return boto3.session.Session(
aws_access_key_id=AwsService.access_key_id,
aws_secret_access_key=AwsService.secret_access_key)
@staticmethod
def get_regions():
return AwsService.get_session().get_available_regions('ssm')

View File

View File

@ -0,0 +1,20 @@
from common.cmd.cmd_result import CmdResult
__author__ = 'itay.mizeretz'
class AwsCmdResult(CmdResult):
"""
Class representing an AWS command result
"""
def __init__(self, command_info):
super(AwsCmdResult, self).__init__(
self.is_successful(command_info), command_info[u'ResponseCode'], command_info[u'StandardOutputContent'],
command_info[u'StandardErrorContent'])
self.command_info = command_info
@staticmethod
def is_successful(command_info):
return command_info[u'Status'] == u'Success'

View File

@ -0,0 +1,39 @@
import time
from common.cloud.aws_service import AwsService
from common.cmd.aws_cmd_result import AwsCmdResult
__author__ = 'itay.mizeretz'
class AwsCmdRunner(object):
"""
Class for running a command on a remote AWS machine
"""
def __init__(self, instance_id, region, is_powershell=False):
self.instance_id = instance_id
self.region = region
self.is_powershell = is_powershell
self.ssm = AwsService.get_client('ssm', region)
def run_command(self, command, timeout):
command_id = self._send_command(command)
init_time = time.time()
curr_time = init_time
command_info = None
while curr_time - init_time < timeout:
command_info = self.ssm.get_command_invocation(CommandId=command_id, InstanceId=self.instance_id)
if AwsCmdResult.is_successful(command_info):
break
else:
time.sleep(0.5)
curr_time = time.time()
return AwsCmdResult(command_info)
def _send_command(self, command):
doc_name = "AWS-RunPowerShellScript" if self.is_powershell else "AWS-RunShellScript"
command_res = self.ssm.send_command(DocumentName=doc_name, Parameters={'commands': [command]},
InstanceIds=[self.instance_id])
return command_res['Command']['CommandId']

View File

@ -0,0 +1,12 @@
class CmdResult(object):
"""
Class representing a command result
"""
def __init__(self, is_success, status_code=None, stdout=None, stderr=None):
self.is_success = is_success
self.status_code = status_code
self.stdout = stdout
self.stderr = stderr

View File

@ -0,0 +1,22 @@
from abc import abstractmethod
__author__ = 'itay.mizeretz'
class CmdRunner(object):
"""
Interface for running a command on a remote machine
"""
# Default command timeout in seconds
DEFAULT_TIMEOUT = 5
@abstractmethod
def run_command(self, command, timeout=DEFAULT_TIMEOUT):
"""
Runs the given command on the remote machine
:param command: The command to run
:param timeout: Timeout in seconds for command.
:return: Command result
"""
raise NotImplementedError()

View File

@ -1,6 +1,6 @@
import logging import logging
from common.cloud.aws import AWS from common.cloud.aws_instance import AwsInstance
__author__ = 'itay.mizeretz' __author__ = 'itay.mizeretz'
@ -15,7 +15,7 @@ class AwsCollector(object):
@staticmethod @staticmethod
def get_aws_info(): def get_aws_info():
LOG.info("Collecting AWS info") LOG.info("Collecting AWS info")
aws = AWS() aws = AwsInstance()
info = {} info = {}
if aws.is_aws_instance(): if aws.is_aws_instance():
LOG.info("Machine is an AWS instance") LOG.info("Machine is an AWS instance")

View File

@ -1,6 +1,6 @@
import cc.auth import cc.auth
from cc.environment import Environment from cc.environment import Environment
from common.cloud.aws import AWS from common.cloud.aws_instance import AwsInstance
__author__ = 'itay.mizeretz' __author__ = 'itay.mizeretz'
@ -8,7 +8,7 @@ __author__ = 'itay.mizeretz'
class AwsEnvironment(Environment): class AwsEnvironment(Environment):
def __init__(self): def __init__(self):
super(AwsEnvironment, self).__init__() super(AwsEnvironment, self).__init__()
self.aws_info = AWS() self.aws_info = AwsInstance()
self._instance_id = self._get_instance_id() self._instance_id = self._get_instance_id()
self.region = self._get_region() self.region = self._get_region()

View File

@ -7,7 +7,7 @@ from botocore.exceptions import UnknownServiceError
from cc.resources.exporter import Exporter from cc.resources.exporter import Exporter
from cc.services.config import ConfigService from cc.services.config import ConfigService
from cc.environment.environment import load_server_configuration_from_file from cc.environment.environment import load_server_configuration_from_file
from common.cloud.aws import AWS from common.cloud.aws_instance import AwsInstance
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -20,7 +20,7 @@ class AWSExporter(Exporter):
@staticmethod @staticmethod
def handle_report(report_json): def handle_report(report_json):
aws = AWS() aws = AwsInstance()
findings_list = [] findings_list = []
issues_list = report_json['recommendations']['issues'] issues_list = report_json['recommendations']['issues']
if not issues_list: if not issues_list:

View File

@ -0,0 +1,20 @@
import json
from flask import request, jsonify, make_response
import flask_restful
class RemoteRun(flask_restful.Resource):
def run_aws_monkey(self, request_body):
instance_id = request_body.get('instance_id')
region = request_body.get('region')
os = request_body.get('os') # TODO: consider getting this from instance
island_ip = request_body.get('island_ip') # TODO: Consider getting this another way. Not easy to determine target interface
def post(self):
body = json.loads(request.data)
if body.get('type') == 'aws':
local_run = self.run_aws_monkey(body)
return jsonify(is_running=local_run[0], error_text=local_run[1])
# default action
return make_response({'error': 'Invalid action'}, 500)