Skip to content

Commit

Permalink
PR: PEP8 Compliance (#7)
Browse files Browse the repository at this point in the history 
* PEP8: Addressed W191 (Indentation with tabs) from pycodestyle

* More PEP8 work, pycodestyle should have no complaints other than line length at this point
  • Loading branch information
@ncc-erik-steringer
ncc-erik-steringer committed Aug 18, 2018
1 parent ce28473 commit c9e13e3
Show file tree
Hide file tree
Showing 15 changed files with 1,166 additions and 1,134 deletions.
280 changes: 142 additions & 138 deletions pmapper.py
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,5 @@
#!/usr/bin/env python

"""
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
Expand All @@ -14,159 +8,169 @@
import botocore.session
import logging
import os.path
import sys

import principalmap.enumerator
from principalmap.querying import perform_query
from principalmap.visualizing import perform_visualization
import sys

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


def main():
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 parsed.picked_cmd == 'graph':
handle_graph(parsed)
elif parsed.picked_cmd == 'query':
handle_query(parsed)
elif parsed.picked_cmd == 'visualize':
handle_visualize(parsed)
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 parsed.picked_cmd == 'graph':
handle_graph(parsed)
elif parsed.picked_cmd == 'query':
handle_query(parsed)
elif parsed.picked_cmd == 'visualize':
handle_visualize(parsed)


def handle_graph(parsed):
if not parsed.display:
graph = pull_graph(parsed.profile)
print('Created an ' + str(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")
graph.write_to_fd(graphfile)
else:
filepath = os.path.join(os.path.expanduser('~'), '.principalmap/graphfile-' + parsed.profile)
graph = graph_from_file(filepath)
print(str(graph))
if not parsed.display:
graph = pull_graph(parsed.profile)
print('Created an ' + str(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")
graph.write_to_fd(graphfile)
else:
filepath = os.path.join(os.path.expanduser('~'), '.principalmap/graphfile-' + parsed.profile)
graph = graph_from_file(filepath)
print(str(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)
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)
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)


def pull_graph(profilearg):
botocore_session = botocore.session.Session(profile=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)
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'])
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()
enumerator = principalmap.enumerator.Enumerator(botocore_session)
enumerator.fillOutGraph()

return enumerator.graph

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
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()
main()
Loading

0 comments on commit c9e13e3

Please sign in to comment.