2
\$\begingroup\$

We have a requirement that access to an AWS environment is enabled on demand, and said access disabled afterwards. In each case (up / down), this means we will do two things:

  1. Enable access

    1. Turn on the Bastion instance
    2. Modify the Security Groups of the EC2 instances, so a SG allowing RDP access is added to the list of associated SGs.

  2. Disable access

    1. Turn off the Bastion
    2. Modify the SGs of the EC2 instances, to remove the SG that grants RDP access.

The script is pretty straight forwards, but there is not much error handling / logging going on.

The script is designed to run either from local machine or in AWS lambda, without modification (providing environment variables, credentials/permissions setup correctly). 

import boto3, os 

region = os.environ['PFregion'] 
security_group_id = os.environ['PFSG'] 
bastion_iid = os.environ['PFbastion_iid'] 
aws_profile = 'profile_name'

def lambda_handler(event, context):
    if (__name__ == "__main__"):
        print 'using ENV creds'     
        session = boto3.Session(profile_name=aws_profile)
        ec2_client = session.client('ec2', region_name=region)      
        ec2_resx = session.resource('ec2', region_name=region)
    else:
        ec2_client = boto3.client('ec2', region_name=region)
        ec2_resx = boto3.resource('ec2', region_name=region)

    all_instances = ec2_resx.instances.filter()

    if (event['up'].lower() == 'true'):
    # Start the Bastion 
        print 'inside up'
        ec2_client.start_instances(
            InstanceIds = [bastion_iid,],
            AdditionalInfo = 'Called from Lambda function to Up access',
            #DryRun = event['dry_run'], 
        )
        # Modify the SGs accordingly 
        for instance in all_instances:
            all_sg_ids = [sg['GroupId'] for sg in instance.security_groups]
            if security_group_id not in all_sg_ids:
                print 'SG not assigned currently'
                all_sg_ids.append(security_group_id)
                instance.modify_attribute(Groups=all_sg_ids)

    else:
        # Stop the Bastion 
        print 'inside down'
        ec2_client.stop_instances(
            InstanceIds = [bastion_iid,],           
            #DryRun = bool(event['dry_run']), 
        )
        # Remove the SG rules
        for instance in all_instances:
            all_sg_ids = [sg['GroupId'] for sg in instance.security_groups]
            if security_group_id in all_sg_ids:
                print 'SG IS assigned currently'
                all_sg_ids.remove(security_group_id)
                instance.modify_attribute(Groups=all_sg_ids)

if __name__ == '__main__':
    event = { "up" : "true", "desc" : "this is a description field", "dry_run" : "false" }
    #print event
    lambda_handler(event, None)
\$\endgroup\$

1 Answer 1

Reset to default
2
\$\begingroup\$

A recommendation that I give with most of my reviews, and which is relevant here, is a perusal of PEP 8, the official style guide for Python. Many of my comments below will reference that document.

import boto3, os

PEP 8 has a few things to say about this (in the section entitled Imports). First, generally you shouldn't put multiple modules on the same line. Next, third-party modules should be distinguished from built-ins, preferably with a blank line (and local application modules after). In addition, I like to alphabetize my imports within each group. It can prevent duplicates. So I would do it like this:

import os

import boto3

region = os.environ['PFregion']
security_group_id = os.environ['PFSG']
bastion_iid = os.environ['PFbastion_iid']
aws_profile = 'profile_name'

I have several problems with this. For one thing, os has a shortcut function for accessing environment variables: os.getenv. I would also pass these as arguments to lambda_handler instead of using global variables. This gives the option of calling the function from another module or calling it multiple times without the need of setting environment variables.

The if __name__ == '__main__': block should be where those variables are gotten. As a nice touch, you could use argparse to add commandline arguments as an alternative to environment variables. I would probably do it something like this:

# At the top with your other imports
from argparse import ArgumentParser

# In `if __name == '__main__':`
parser = ArgumentParser(description="Your program description")
parser.add_argument('-r', '--region', description='PF region')
 # (Repeat for other arguments)

args = parser.parse_args()

Then you can pass something like args.region or os.getenv('PFregion') to your function. That gives the command line option as the decision-maker, but it falls back on the environment variable.

def lambda_handler(event, context):
    if (__name__ == "__main__"):

This is a worker function; it shouldn't care what is calling it. I would move the block outside the function and pass the relevant variables as arguments. You could even use something like:

# Though because of my previous recommendation, this function would have more arguments.
def lambda_handler(event, context, session=boto3):
    ec2_client = session.client('ec2', region_name=region)
    ec2_resx = session.resource('ec2', region_name=region)

Then your if __name__ == '__main__': block could simply pass the session as an argument, but the default without a session still works.

event = { "up" : "true", "desc" : "this is a description field", "dry_run" : "false" }

I would add these options to the environment variable and command line options.

Python has its own booleans. Use them: "up": True.

Another styling comment: I wouldn't put spaces on both sides of the braces. PEP 8 comments on this.

I also noticed that you never use event['desc'] or event['dry_run']. Why then are they ever defined?

\$\endgroup\$
0

You must log in to answer this question.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.