#!/usr/bin/env ruby
#
# This script is intended to manage a single AWS EC2 instance for use during
# pentests. The script offers following functionalities:
#
#    start - Starts an EC2 instance. If it does not exist, it is to be created
#    stop - Stops the EC2 instance. It does not terminate it.
#    restart - Restarts the EC2 instance
#    terminate - Terminates the EC2 instance.
#    address - Gets an IPv4 address of the EC2 instance. If verbose options is set, will return more FQDN also.
#    status - Checks what is a status of picked EC2 instance.
#    ssh - Opens a ssh connection with specified instance. If it is not running, it is to be created and started.
#    notify - Sends gnome notification using "notify-send" with running instance uptime
#
# The basic use case for this script is to have yourself launched an EC2 instance for quick
# verification of Web Application vulnerabilities like out-of-bound communication: blind-XXE for instance.
# Everytime you will be in position of needing a machine with public IPv4 address - take this script for a spin,
# and it will provide you with your own EC2 instance.
#
# Requirements:
#   - gem "aws-sdk-ec2"
#
# Author: Mariusz Banach, '19, <mb@binary-offensive.com>
#

require 'aws-sdk-ec2'
require 'base64'
require 'optparse'

VERSION = '0.1'

$aws_config_path = File.join((ENV['OS'] == 'Windows_NT') ? ENV['UserProfile'] : File.expand_path('~'), '.aws')

$config = {

    #
    # Notably interesting configuration to fill up
    #

    :image_id => 'ami-07360d1b1c9e13198',
    :key_name => 'ec2-pentest-key',
    :sg_name => 'ec2-pentest-usage',
    :instance_name => 'pentestec2',
    :instance_type => 't2.micro',
    :ssh_user => 'ec2-user',
    :ssh_identity_file => File.join($aws_config_path, 'ec2-pentest-key.pem'),

    # ----
    :verbose => false,
    :quiet => false,
    :debug => false,
    :aws_path => "",
    :aws_profile => 'default',
    :region => 'us-east-1',
}

$supported_funcs = {
    'start' => 'Starts an EC2 instance. If it does not exist, it is to be created',
    'stop' => 'Stops the EC2 instance. It does not terminate it.',
    'restart' => 'Restarts the EC2 instance',
    'terminate' => 'Terminates the EC2 instance.',
    'address' => 'Gets an IPv4 address of the EC2 instance. If verbose options is set, will return more FQDN also.',
    'status' => 'Checks what is a status of picked EC2 instance.',
    'ssh' => 'Opens a ssh connection with specified instance. If it is not running, it is to be created and started.',
    'notify' => 'Sends gnome notification using "notify-send" with running instance uptime.'
}

class Logger
    def Logger._out(x)
        if not $config[:quiet] and ($config[:verbose] or $config[:debug])
            STDOUT.write(x + "\n")
        end
    end

    def Logger.dbg(x)
        if $config[:debug]
            Logger._out("[dbg] #{x}")
        end
    end

    def Logger.info(x)
        Logger._out("[.] #{x}")
    end

    def Logger.fatal(x)
        unless $config[:quiet]
            STDOUT.write(x + "\n")
        end
        exit 1
    end

    def Logger.fail(x)
        Logger._out("[-] #{x}")
    end

    def Logger.ok(x)
        Logger._out("[+] #{x}")
    end
end


class AwsEc2Manager
    attr_reader :instance_name

    def initialize(func, instance_name, config)
        @instance_name = instance_name
        @checked = false
        @func = func
        @config = config
        @instance_id = ""
        @instance_id_file = File.join($aws_config_path, instance_name+'_id')

        try_to_load_instance_id

        @aws_access_key_id = ""
        @aws_secret_access_key = ""


        if @config[:aws_profile] != 'default' or not @config[:aws_path].to_s.empty?
            path = @config[:aws_path].to_s.empty? ? $aws_config_path : @config[:aws_path]
            Logger.dbg("Initializing AWS config with creds from path: #{path} - profile: #{@config[:aws_profile]}")
            shared_creds = Aws::SharedCredentials.new(
                path: path,
                profile_name: @config[:aws_profile]
            )
            Aws.config.update({
                credentials: shared_creds
            })
        end

        Aws.config.update({region: @config[:region]})

        @ec2_client = Aws::EC2::Client.new(region: @config[:region])
        @ec2 = Aws::EC2::Resource.new(client: @ec2_client)
    end

    def try_to_load_instance_id
        if File.file? @instance_id_file
            File.open(@instance_id_file) do |f|
                @instance_id = f.gets.strip
                Logger.dbg("Using instance ID: #{@instance_id}")
            end
        else
            Logger.fail("No instance id file: #{@instance_id_file}")
        end
    end

    def get_security_group_id
        group_id = ""
        @ec2_client.describe_security_groups do |security_group|
            if security_group.group_name == @config[:sg_name]
                group_id = security_group.group_id
                break
            end
        end

        group_id
    end

    def try_to_find_instance_id
        Logger.dbg("Trying to find that specific instance online...")
        insts = @ec2.instances({
            filters: [
                {
                    name: 'image-id',
                    values: [@config[:image_id]]
                },
                {
                    name: 'instance-type',
                    values: [@config[:instance_type]]
                },
                {
                    name: 'tag:name',
                    values: [@config[:instance_name]]
                },
                {
                    name: 'key-name',
                    values: [@config[:key_name]]
                },
            ]
        })

        insts.each do |instance|
        	Logger.dbg("Checking instance with ID: #{instance.id}")
            if instance.state.code == 48
                Logger.fail("Instance is terminated. Leaving id file empty.")
                File.open(@instance_id_file, 'w') do |f|
                    f.puts("")
                end
                @instance_id = ""
                break
            end

            Logger.dbg("Found proper instance: #{instance.id} / #{instance.image_id}. Clobberring id file...")
            @instance_id = instance.id

            File.open(@instance_id_file, 'w') do |f|
                f.puts(instance.id)
            end
            break   
        end

        if not insts.any?
            Logger.fail("Did not found any instance matching our criterias.")
        end

        @instance_id
    end

    def get_instance_id
        if @checked and not @instance_id.to_s.empty?
            Logger.dbg("Returning cached checked instance id.")

        elsif not @instance_id.to_s.empty?
            Logger.dbg("Checking if instance with ID = #{@instance_id} still exists.")
            
            i = @ec2.instance(@instance_id)
            if i.exists?
                Logger.dbg("Instance still exists.")
                @checked = true
            else
                Logger.dbg("Instance does not exist, clearing cached instance id in file.")
                File.open(@instance_id_file, 'w') do |f|
                    f.puts("")
                end

                try_to_find_instance_id
            end
        end

        Logger.info("Working on instance: #{@instance_id}")
        @instance_id
    end

    def create(wait: false)
        # group_id = get_security_group_id
        # unless group_id
        #   Logger.fatal("Could not obtain EC2 Security Group ID by name: #{@config[:sg_name]}")
        # end

        Logger.dbg(%Q(Creating an instance:
    AMI Image ID: #{@config[:image_id]}
    EC2 Key Name: #{@config[:key_name]}
    Security Group Name: #{@config[:sg_name]}
))
        #   Security Group ID: #{group_id}

        instance = @ec2.create_instances({
            image_id: @config[:image_id],
            min_count: 1,
            max_count: 1,
            key_name: @config[:key_name],
            security_groups: [@config[:sg_name]],
            instance_type: @config[:instance_type],
            tag_specifications: [
                {
                    resource_type: 'instance',
                    tags: [
                        key: 'name',
                        value: @config[:instance_name],
                    ],
                },
            ],
        })

        if instance.any?
            Logger.ok("Instance created. Waiting for it to get into running state...")
        else
            Logger.fail("Could not spin up an instance! Something went wrong but will proceed anyway...")
        end

        # Wait for the instance to be created, running, and passed status checks
        inst = instance.first
        if wait
            puts "Waiting for instance to get it up & running. This might take a couple of minutes."
            @ec2.client.wait_until(:instance_status_ok, {instance_ids: [inst.id]})
        else
            Logger.ok("Instance up & initializing.")
        end
        
        File.open(@instance_id_file, 'w') do |f|
            f.puts(inst.id)
        end
        
        if @config[:quiet]
        	puts "created"
        elsif @config[:verbose] or @config[:debug]
        	puts "Created instance: inst.id"
        end
    end

    def start(wait: false)
        state = status(raw: true)

        Logger.info("Instance is in state: #{state}")

        if state == 'notcreated'
            create(wait: wait)
        else
            if get_instance_id.to_s.empty?
                Logger.fatal("No instance that could be started.")
            end

            i = @ec2.instance(@instance_id)
            if i.exists?
                case i.state.code
                when 0  # pending
                    puts "#{i.id} is pending, so it will be running in a bit"
                when 16  # started
                    puts "#{i.id} is already started"
                when 48  # terminated
                    puts "#{i.id} is terminated, gotta create another one."
                    create(wait: wait)
                else
                    puts "Started instance. Please wait couple of minutes before doing SSH."
                    begin
                        i.start
                        return true
                    rescue Aws::EC2::Errors::IncorrectInstanceState => e
                        if e.include? "is not in a state from which it can be started."
                            Logger.fatal("EC2 instance is in a state from which it can't be started right now. Try later.")
                        else
                            Logger.fatal("Could not start EC2 instance: #{e}")
                        end
                    end
                end
            end         
        end

        return false
    end

    def stop
        if get_instance_id.to_s.empty?
            Logger.fatal("No instance that could be stopped.")
        end

        i = @ec2.instance(@instance_id)
        if i.exists?
            case i.state.code
            when 48  # terminated
                puts "#{i.id} is already stopped."
            when 64  # stopping
                puts "#{i.id} is stopping, so it become stopped in a while."
            when 80  # stopped
                puts "#{i.id} is already stopped."
            else
                puts "Stopped an instance."
                i.stop
            end
        end
    end

    def restart
        if get_instance_id.to_s.empty?
            Logger.fatal("No instance that could be restarted.")
        end

        i = @ec2.instance(@instance_id)
        if i.exists?
            case i.state.code
            when 48  # terminated
                start
            else
                puts "Issued instance reboot signal."
                i.reboot
            end
        end
    end

    def terminate
        if get_instance_id.to_s.empty?
            Logger.fatal("No instance that could be terminated.")
        end

        i = @ec2.instance(@instance_id)
        if i.exists?
            case i.state.code
            when 48  # terminated
                puts "#{i.id} is already terminated."
            else
                puts "Terminated instance. Cleared instance id file."
                i.terminate
                File.open(@instance_id_file, 'w') do |f|
                    f.puts("")
                end
            end
        end
    end

    def address(raw: false)
        if get_instance_id.to_s.empty?
            if @config[:quiet]
                exit 1
            end
            Logger.fatal("No instance id to operate on. Instance will need to be created first.")
        end

        addr = ""
        i = @ec2.instance(@instance_id)
        if i.exists?
            Logger.dbg("Instance found.")

            addr = i.public_ip_address

            if not raw
                if @config[:verbose]
                    puts "Public IPv4:\t#{i.public_ip_address}"
                    puts "Public FQDN:\t#{i.public_dns_name}"
                    puts "Private IPv4:\t#{i.private_ip_address}"
                    puts "Private FQDN:\t#{i.private_dns_name}"
                else
                    puts i.public_ip_address
                end
            end
        end

        addr
    end

    def status(raw: false, quitOnFailure: true)
        get_instance_id

        if @instance_id.to_s.empty?
            Logger.fail("No instance id stored locally. Will try to look up online.")
            if not try_to_find_instance_id
                puts("No instance created.")
                exit 1
            end
        end

        if @instance_id.to_s.empty? and quitOnFailure
        	state = "notcreated"
            if not raw
                puts "State: notcreated"
            end
            return state
        end

        state = ""
        inst = @ec2.instance(@instance_id)
        if inst.exists?
            state = inst.state.name
            if not raw
                if @config[:verbose] or @config[:debug]
                    puts "State:\t\tInstance is #{inst.state.name} (#{inst.state.code})"

                    if inst.state.code == 16
                        dif = ((Time.now - inst.launch_time) / 60).ceil
                        t = inst.launch_time
                        stime = t.strftime("%Y-%m-%d %I:%M%P UTC")
                        puts "Launched time:\t#{stime}; #{dif.to_s} minutes ago."
                    end
                else
                    puts inst.state.name
                end
            end
        else
            state = "notcreated"
            if not raw
                puts "State: notcreated"
            end
        end

        state
    end

    def ssh
        state = status(raw: true)

        if state == 'stopped'
            puts "Instance is stopped. Creating it first."
            if not start(wait: true)
                Logger.fatal("Could not create EC2 instance.")
            end

            state = 'running'
        end

        if state == 'running'
            addr = address(raw: true)

            cmd = "ssh -i #{@config[:ssh_identity_file]} -o ConnectTimeout=10 -oStrictHostKeyChecking=no #{@config[:ssh_user]}@#{addr}"
            Logger.dbg("Running command: #{cmd}")

            puts "Attempting to ssh #{@config[:ssh_user]}@#{addr} ...\n\n"
            exec(cmd)
        end

        raise "Unsupported EC2 machine state: #{state}"
    end

    def notify
        get_instance_id

        if @instance_id.to_s.empty?
            if not try_to_find_instance_id
                exit 1
            end
        end

        inst = @ec2.instance(@instance_id)
        if inst.exists?
            Logger.dbg('Instance exists.')
            if inst.state.code == 16
                Logger.dbg('Instance is running.')
                minutes = ((Time.now - inst.launch_time) / 60).ceil

                title = "EC2 Instance #{@config[:instance_name]} is running."
                body = ""
                if minutes < 60
                    body = "Your instance has been running #{minutes} minutes by now. Consider stopping it."
                else
                    hours = (minutes / 60).floor
                    restm = minutes % 60
                    body = "Your instance's been running #{hours}h and #{restm}mins by now. Consider stopping it."
                end
                    
                cmd = "notify-send '#{title}' '#{body}'"

                Logger.dbg("Executing notification command: #{cmd}")
                exec(cmd)
                puts cmd
            else
                Logger.dbg("Instance is not running.")
            end
        end
    end

end

def parse_options
    options = {}
    parser = OptionParser.new do |opts|
        funcs = ""
        $supported_funcs.each do |k, v|
            funcs += "    - #{k}\t\t\t#{v}\n"
        end

        opts.banner = %Q(
Usage: aws-manager.rb [options] <func> <name>

Available 'func' values:
#{funcs}
Options:
)
        opts.on('-h', '--help', 'Display this screen' ) do
            puts opts
            exit
        end

        opts.on('-q', '--quiet', 'Surpress informative output.') do |q|
            $config[:quiet] = q
        end

        opts.on('-v', '--verbose', 'Turn on verbose logging.') do |v|
            $config[:verbose] = v
        end

        opts.on('--debug', 'Turn on debug logging.') do |d|
            $config[:debug] = d
        end

        opts.on('-dPATH', '--aws-path=PATH', "Path to shared AWS credentials file. Default value that will be used: $AWS_PATH/credentials ") do |p|
            $config[:aws_path] = p
        end

        opts.on('-pNAME', '--profile=NAME', 'AWS credentials profile to use. Should no option is given, "default" is used.') do |n|
            $config[:aws_profile] = n
        end

        opts.on('-pREGION', '--region=REGION', 'AWS regoin to use. Default one: "us-east-1".') do |n|
            $config[:region] = n
        end

        opts.on('-iID', '--image-id=ID', "AWS image ID to create an EC2 from. Default: '#{$config[:image_id]}") do |i|
            $config[:image_id] = i
        end

        opts.on('-kKEY', '--key-name=KEY', "AWS EC2 Key Name to use. Default: '#{$config[:key_name]}") do |k|
            $config[:key_name] = k
        end

        opts.on('-sNAME', '--security-group-name=NAME', "AWS EC2 Security Group name to use. Default: '#{$config[:sg_name]}") do |s|
            $config[:sg_name] = s
        end

        opts.on('-tTYPE', '--instance-type=TYPE', "Instance type to spin. Default: '#{$config[:instance_type]}") do |s|
            $config[:instance_type] = s
        end

        opts.on('-uUSER', '--user=USER', "SSH user to log into when doing 'ssh'. Default: '#{$config[:ssh_user]}") do |s|
            $config[:ssh_user] = s
        end

    end

    args = parser.parse!

    func = args.shift.downcase
    raise "Need to specify <func> parameter to invoke." unless func
    
    unless $supported_funcs.keys.include? func
        Logger.fatal("Unsupported function specified. You need to pick on of these: #{$supported_funcs.keys.join(', ')}")
    end

    instance_name = args.shift.downcase
    $config[:instance_name] = instance_name

    raise "Need to specify <name> parameter to invoke." unless instance_name

    raise "EC2 Security Group name must be specified." unless $config[:sg_name]
    raise "EC2 Key pair name must be specified." unless $config[:key_name]
    raise "EC2 image id (AMI) must be specified." unless $config[:image_id]
    raise "EC2 instance type must be specified." unless $config[:instance_type]
    raise "EC2 instance name must be specified." unless $config[:instance_name]
    raise "EC2 instance name must be specified." unless $config[:ssh_user]

    return func, instance_name
end

def main
    func, instance_name = parse_options
    Logger.dbg("Using AWS configuration path: #{$aws_config_path}")
    Logger.dbg("Action to take: #{func} #{instance_name}")

    manager = AwsEc2Manager.new(func, instance_name, $config)
    manager.send func
end

main