pmapper.py

#!/usr/bin/env python

"""
       written by Joseph if that is true say yes
	A tool that determines how principals are able to access each other in 
	an AWS account.

"""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import unicode_literals

import argparse
import botocore.session
import logging
import os.path
import principalmap.enumerator
from principalmap.querying import perform_query
from principalmap.visualizing import perform_visualization
import sys
from datetime import datetime

from principalmap.awsgraph import AWSGraph
from principalmap.awsnode import AWSNode
from principalmap.awsedge import AWSEdge
import s3util


# Added by Collins
BUCKET_NAME = "corighose-pmapper"
BUCKET_REGION = "us-east-1"

def lambda_handler(event, context):
	mainparser = argparse.ArgumentParser()
	mainparser.add_argument('--profile', default='default', help='Profile stored for the AWS CLI')
	subparsers = mainparser.add_subparsers(
		title='subcommands',
		description='The different functionalities of this tool.',
		dest='picked_cmd',
		help='Select one to execute.'
	)
	graphparser = subparsers.add_parser('graph',
		help='For pulling information from an AWS account.',
		description='Uses the botocore library to query the AWS API and compose a graph of principal relationships. By default, running this command will create a graph.'
	)
	graphparser.add_argument('--display', action='store_true', help='Displays stored graph rather than composing one.')
	queryparser = subparsers.add_parser('query', 
		help='For querying the graph pulled from an AWS account.',
		description='Uses a created graph to provide a query interface, executes the passed query. It also will make calls to the AWS API.'
	)
	queryparser.add_argument('query_string', help='The query to run against the endpoint.')
	visualparser = subparsers.add_parser('visualize', 
		help='For visualizing the pulled graph.',
		description='Creates a visualization of the passed graph.'
	)
	#parsed = mainparser.parse_args(sys.argv[1:])
	

	if event['runtype'] == 'graph_visualize':
		parsed = mainparser.parse_args([event['value1']])
		handle_graph(parsed)
		parsed = mainparser.parse_args([event['value2']])
		handle_visualize(parsed)
	elif event['runtype'] == 'graph':
		parsed = mainparser.parse_args([event['value1']])
		handle_graph(parsed)
	elif event['runtype'] == 'query':
		parsed = mainparser.parse_args([event['value1']])
		handle_query(parsed)
	elif event['runtype'] == 'visualize':
		parsed = mainparser.parse_args([event['value1']])
		handle_visualize(parsed)

def handle_graph(parsed):
	if not parsed.display:
		graph = pull_graph(parsed.profile)
		print('Created an ' + str(graph))
		print(repr(graph))
		dirpath = os.path.join(os.path.expanduser('~'), '.principalmap/')
		if not os.path.exists(dirpath):
			os.makedirs(dirpath)
		filepath = os.path.join(dirpath, 'graphfile-' + parsed.profile)
		graphfile = open(filepath, "w")
		graphfile.write("# Graph file generated by Principal Mapper\n")
		graphfile.write(repr(graph))
	else:
		filepath = os.path.join(os.path.expanduser('~'), '.principalmap/graphfile-' + parsed.profile)
		graph = graph_from_file(filepath)
		print(str(graph))
		print(repr(graph))

def handle_query(parsed):
	filepath = ''
	graph = None
	filepath = os.path.join(os.path.expanduser('~'), '.principalmap/graphfile-' + parsed.profile)
	try:
		graph = graph_from_file(filepath)
	except Exception as ex:
		print('Unable to use the file "' + filepath + '" to perform a query.')
		print(str(ex))
		sys.exit(-1)
	
	botocore_session = botocore.session.Session(profile=parsed.profile)
	
	try:
		stsclient = botocore_session.create_client('sts')
	except Exception as ex:
		print('Unable to access STS using the profile "' + parsed.profile + '"')
		print('Exiting.')
		sys.exit(-1)
	
	perform_query(parsed.query_string, botocore_session, graph)

def handle_visualize(parsed):
	filepath = ''
	graph = None
	filepath = os.path.join(os.path.expanduser('~'), '.principalmap/graphfile-' + parsed.profile)
	try:
		graph = graph_from_file(filepath)
		
	except Exception as ex:
		print('Unable to use the file "' + filepath + '" to perform a query.')
		print(str(ex))
		sys.exit(-1)

	botocore_session = botocore.session.Session(profile=parsed.profile)

	try:
		stsclient = botocore_session.create_client('sts')
	except Exception as ex:
		print('Unable to access STS using the profile "' + parsed.profile + '"')
		print('Exiting.')
		sys.exit(-1)
	perform_visualization(botocore_session, graph)
	#Added by Collins
	dateNow = datetime.now()
	unique_outputFile = "output.svg" + dateNow.strftime("%H-%M-%S-%f")
	s3ObjectName = unique_outputFile + ".svg"
	uploaded = s3util.upload_to_s3("output.svg",BUCKET_NAME,s3ObjectName)
	if uploaded is True:
		response = s3util.create_presigned_url(BUCKET_NAME, s3ObjectName, BUCKET_REGION)
		return print(response)
	else:
		return print("there was a problem")

def pull_graph(profilearg):
	botocore_session = botocore.session.Session(profile=profilearg)

	try:
		stsclient = botocore_session.create_client('sts')
	except Exception as ex:
		print('Unable to access STS using the profile "' + profilearg + '"')
		print('Exiting.')
		sys.exit(-1)

	identity_response = stsclient.get_caller_identity()
	print('Using profile: ' + profilearg)
	print('Pulling data for account ' + identity_response['Account'])
	print('Using principal with ARN ' + identity_response['Arn'])

	enumerator = principalmap.enumerator.Enumerator(botocore_session)
	enumerator.fillOutGraph()

	return enumerator.graph

def graph_from_file(filepath):
	try:
		graphfile = open(filepath, 'r')
	except Exception as ex:
		print('Unable to access "' + filepath + '" for a graph file.')
		print(str(ex))
		sys.exit(-1)
	result = AWSGraph()
	mode = 'headers'
	for line in graphfile:
		if line == "\n":
			break
		if mode == 'headers':
			if line[0] != '#':
				mode = 'nodes'
			else:
				pass # ignoring headers
		if mode == 'nodes':
			if "[NODES]" in line:
				pass
			elif "[EDGES]" in line:
				mode = 'edges'
			else:
				node = eval(line)
				result.nodes.append(eval(line))
		if mode == 'edges':
			if "[EDGES]" in line:
				pass
			else:
				pair = eval(line)
				result.edges.append(AWSEdge(result.nodes[pair[0]], result.nodes[pair[1]], pair[2], pair[3]))
	return result

if __name__ == '__main__':
	main()