Skip to content

Commit 8126223

Browse files
mvanholsteijnmvanholsteijn
committed
working neo4j importer
1 parent f369887 commit 8126223

15 files changed

+386
-3
lines changed

‎aws_visualizer/__init__.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
from .graph_region import main
1+
from aws_visualizer.dot.graph_region import main

‎aws_visualizer/dot/__init__.py

Whitespace-only changes.

‎aws_visualizer/graph_region.py ‎aws_visualizer/dot/graph_region.py

File renamed without changes.

‎aws_visualizer/neo4j/__init__.py

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
from aws_visualizer.neo4j.main import main

‎aws_visualizer/neo4j/ec2_instance.py

+48
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
import os
2+
import boto3
3+
from neomodel import *
4+
import logging
5+
from .security_group import SecurityGroup
6+
7+
8+
class EC2Instance(StructuredNode):
9+
Name = StringProperty()
10+
ARN = StringProperty(unique_index=True)
11+
OwnerId = StringProperty()
12+
InstanceId = StringProperty()
13+
SecurityGroups = Relationship('SecurityGroup', 'APPLIED')
14+
VpcId = StringProperty()
15+
16+
@staticmethod
17+
def arn(region, owner_id, instance_id):
18+
return 'arn::ec2:{}:{}:instance/{}'.format(region, owner_id, instance_id)
19+
20+
def get_instance_name(obj):
21+
return next(map(lambda t: t['Value'], list(filter(lambda t: t['Key'] == 'Name', obj['Tags'] if 'Tags' in obj else {}))), obj['InstanceId'])
22+
23+
def get_account_id_and_region():
24+
sts = boto3.client('sts')
25+
return sts.get_caller_identity()['Account'], sts.meta.region_name
26+
27+
def load():
28+
ec2 = boto3.client('ec2')
29+
OwnerId, AWSRegion = get_account_id_and_region()
30+
paginator = ec2.get_paginator('describe_instances')
31+
for page in paginator.paginate():
32+
for r in page['Reservations']:
33+
for i in filter(lambda i: i['State']['Name'] != "terminated", r['Instances']):
34+
InstanceId = i['InstanceId']
35+
arn = EC2Instance.arn(AWSRegion, OwnerId, InstanceId)
36+
name = get_instance_name(i)
37+
instance = EC2Instance.create_or_update({'ARN': arn, 'Name': get_instance_name(i), 'InstanceId': InstanceId, 'OwnerId': OwnerId})[0]
38+
sys.stderr.write('INFO: loading ec2 instance {}\n'.format(InstanceId))
39+
40+
# remove old associations
41+
for group in instance.SecurityGroups:
42+
instance.SecurityGroups.disconnect(group)
43+
44+
for sg in i['SecurityGroups']:
45+
group = SecurityGroup.get_or_create({'ARN': SecurityGroup.arn(ec2.meta.region_name, OwnerId, sg['GroupId'])})[0]
46+
sys.stderr.write('INFO: applying security group {} to ec2 instance {}\n'.format(group.GroupId, InstanceId))
47+
instance.SecurityGroups.connect(group)
48+

‎aws_visualizer/neo4j/elb.py

+50
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
import os
2+
import boto3
3+
from neomodel import *
4+
import logging
5+
6+
from .security_group import SecurityGroup
7+
from .ec2_instance import EC2Instance
8+
from .sts import get_account_id_and_region
9+
10+
class ElasticLoadBalancer(StructuredNode):
11+
Name = StringProperty()
12+
ARN = StringProperty(unique_index=True)
13+
SecurityGroups = Relationship('SecurityGroup', 'APPLIED')
14+
Backends = Relationship('EC2Instance', 'BACKEND')
15+
VpcId = StringProperty()
16+
17+
@staticmethod
18+
def arn(region, owner_id, name):
19+
return 'arn:aws:elasticloadbalancing:{}:{}:loadbalancer/{}'.format(region, owner_id, name)
20+
21+
22+
def load():
23+
elb = boto3.client('elb')
24+
OwnerId, AWSRegion = get_account_id_and_region()
25+
paginator = elb.get_paginator('describe_load_balancers')
26+
for page in paginator.paginate():
27+
for lb in page['LoadBalancerDescriptions']:
28+
name = lb['LoadBalancerName']
29+
arn = ElasticLoadBalancer.arn(AWSRegion, OwnerId, name)
30+
properties = { 'Name': name, 'ARN': arn, 'OwnerId': OwnerId, 'VpcId': lb['VPCId'] if 'VPCId' in lb else None}
31+
loadbalancer = ElasticLoadBalancer.create_or_update(properties)[0]
32+
sys.stderr.write('INFO: loading classic elb {}\n'.format(name))
33+
34+
for group in loadbalancer.SecurityGroups:
35+
loadbalancer.SecurityGroups.disconnect(group)
36+
37+
for instance in loadbalancer.Backends:
38+
loadbalancer.Backends.disconnect(instance)
39+
40+
for GroupId in lb['SecurityGroups']:
41+
group = SecurityGroup.get_or_create({'ARN': SecurityGroup.arn(AWSRegion, OwnerId, GroupId)})[0]
42+
sys.stderr.write('INFO: applying security group {} to elb {}\n'.format(group.GroupId, name))
43+
loadbalancer.SecurityGroups.connect(group)
44+
45+
for InstanceId in map(lambda i: i['InstanceId'], lb['Instances']):
46+
arn = EC2Instance.arn(AWSRegion, OwnerId, InstanceId)
47+
instance = EC2Instance.get_or_create({'ARN': arn, 'InstanceId': InstanceId, 'OwnerId': OwnerId})[0]
48+
sys.stderr.write('INFO: adding instance {} as backed to elb {}\n'.format(InstanceId, name))
49+
loadbalancer.Backends.connect(instance)
50+

‎aws_visualizer/neo4j/elbv2.py

+29
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
import os
2+
import boto3
3+
from neomodel import *
4+
import logging
5+
6+
from .sts import get_account_id_and_region
7+
from .security_group import SecurityGroup
8+
from .elb import ElasticLoadBalancer
9+
10+
11+
def load():
12+
elb = boto3.client('elbv2')
13+
OwnerId, AWSRegion = get_account_id_and_region()
14+
paginator = elb.get_paginator('describe_load_balancers')
15+
for page in paginator.paginate():
16+
for lb in page['LoadBalancers']:
17+
name = lb['LoadBalancerName']
18+
arn = lb['LoadBalancerArn']
19+
properties = { 'Name': name, 'ARN': arn, 'OwnerId': OwnerId, 'VpcId': lb['VpcId']}
20+
loadbalancer = ElasticLoadBalancer.create_or_update(properties)[0]
21+
sys.stderr.write('INFO: loading elb v2 {}\n'.format(name))
22+
23+
for group in loadbalancer.SecurityGroups:
24+
loadbalancer.SecurityGroups.disconnect(group)
25+
26+
for GroupId in lb['SecurityGroups']:
27+
group = SecurityGroup.get_or_create({'ARN': SecurityGroup.arn(AWSRegion, OwnerId, GroupId)})[0]
28+
sys.stderr.write('INFO: applying security group {} to elb v2 {}\n'.format(group.GroupId, name))
29+
loadbalancer.SecurityGroups.connect(group)

‎aws_visualizer/neo4j/function.py

+42
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
import os
2+
import boto3
3+
from neomodel import *
4+
import logging
5+
6+
from .sts import get_account_id_and_region
7+
from .security_group import SecurityGroup
8+
from .elb import ElasticLoadBalancer
9+
10+
11+
class Function(StructuredNode):
12+
Name = StringProperty()
13+
ARN = StringProperty(unique_index=True)
14+
SecurityGroups = RelationshipFrom('SecurityGroup', 'APPLIED')
15+
OwnerId = StringProperty()
16+
VpcId = StringProperty()
17+
18+
@staticmethod
19+
def arn(region, owner_id, name):
20+
return 'arn:aws:elasticloadbalancing:{}:{}:loadbalancer/{}'.format(region, owner_id, name)
21+
22+
def load():
23+
awslambda = boto3.client('lambda')
24+
OwnerId, AWSRegion = get_account_id_and_region()
25+
paginator = awslambda.get_paginator('list_functions')
26+
for page in paginator.paginate():
27+
for f in filter(lambda f: 'VpcConfig' in f and f['VpcConfig']['VpcId'] != '', page['Functions']):
28+
name = f['FunctionName']
29+
arn = f['FunctionArn']
30+
properties = { 'Name': name, 'ARN': arn, 'OwnerId': OwnerId, 'VpcId': f['VpcConfig']['VpcId']}
31+
function = Function.create_or_update(properties)[0]
32+
sys.stderr.write('INFO: loading function {}\n'.format(name))
33+
34+
# remove old associations
35+
for group in function.SecurityGroups:
36+
function.SecurityGroups.disconnect(group)
37+
38+
for GroupId in f['VpcConfig']['SecurityGroupIds']:
39+
group = SecurityGroup.get_or_create({'ARN': SecurityGroup.arn(AWSRegion, OwnerId, GroupId)})[0]
40+
sys.stderr.write('INFO: applying security group {} to function {}\n'.format(group.GroupId, name))
41+
function.SecurityGroups.connect(group)
42+

‎aws_visualizer/neo4j/main.py

+39
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import logging
2+
import os
3+
4+
from neomodel import db, install_all_labels
5+
6+
import aws_visualizer.neo4j.ec2_instance
7+
import aws_visualizer.neo4j.elb
8+
import aws_visualizer.neo4j.elbv2
9+
import aws_visualizer.neo4j.function
10+
import aws_visualizer.neo4j.rds
11+
import aws_visualizer.neo4j.redshift
12+
import aws_visualizer.neo4j.security_group
13+
14+
15+
class AWSNeo4jImporter(object):
16+
def __init__(self):
17+
self.url = os.getenv('NEO4J_BOLT_URL', 'bolt://neo4j:neo4j@localhost:7687')
18+
19+
def login(self):
20+
db.set_connection(self.url)
21+
install_all_labels()
22+
23+
def load(self):
24+
aws_visualizer.neo4j.security_group.load()
25+
aws_visualizer.neo4j.ec2_instance.load()
26+
aws_visualizer.neo4j.elb.load()
27+
aws_visualizer.neo4j.elbv2.load()
28+
aws_visualizer.neo4j.rds.load()
29+
aws_visualizer.neo4j.redshift.load()
30+
aws_visualizer.neo4j.function.load()
31+
32+
def main():
33+
logging.basicConfig(level=os.getenv('LOG_LEVEL', 'ERROR'))
34+
importer = AWSNeo4jImporter()
35+
importer.load()
36+
37+
38+
if __name__ == '__main__':
39+
main()

‎aws_visualizer/neo4j/rds.py

+39
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import os
2+
import boto3
3+
from neomodel import *
4+
import logging
5+
6+
7+
from .sts import get_account_id_and_region
8+
from .security_group import SecurityGroup
9+
10+
class DBInstance(StructuredNode):
11+
DBName = StringProperty()
12+
DBInstanceIdentifier = StringProperty()
13+
ARN = StringProperty(unique_index=True)
14+
SecurityGroups = RelationshipFrom('SecurityGroup', 'APPLIED')
15+
OwnerId = StringProperty()
16+
VpcId = StringProperty()
17+
18+
def load():
19+
rds = boto3.client('rds')
20+
OwnerId, AWSRegion = get_account_id_and_region()
21+
paginator = rds.get_paginator('describe_db_instances')
22+
for page in paginator.paginate():
23+
for db in page['DBInstances']:
24+
name = db['DBInstanceIdentifier']
25+
DBName = db['DBName']
26+
arn = db['DBInstanceArn']
27+
properties = { 'DBName': DBName, 'ARN': arn, 'OwnerId': OwnerId,
28+
'VpcId': db['DBSubnetGroup']['VpcId'], 'DBInstanceIdentifier': name}
29+
db_instance = DBInstance.create_or_update(properties)[0]
30+
sys.stderr.write('INFO: loading rds instance {}\n'.format(name))
31+
32+
# remove old associations
33+
for group in db_instance.SecurityGroups:
34+
db_instance.SecurityGroups.disconnect(group)
35+
36+
for GroupId in map(lambda g: g['VpcSecurityGroupId'], db['VpcSecurityGroups']):
37+
group = SecurityGroup.get_or_create({'ARN': SecurityGroup.arn(AWSRegion, OwnerId, GroupId)})[0]
38+
sys.stderr.write('INFO: applying security group {} to rds {}\n'.format(group.GroupId, name))
39+
db_instance.SecurityGroups.connect(group)

‎aws_visualizer/neo4j/redshift.py

+44
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import os
2+
import boto3
3+
from neomodel import *
4+
import logging
5+
6+
from .sts import get_account_id_and_region
7+
from .security_group import SecurityGroup
8+
9+
10+
class RedshiftCluster(StructuredNode):
11+
DBName = StringProperty()
12+
ClusterIdentifier = StringProperty()
13+
ARN = StringProperty(unique_index=True)
14+
SecurityGroups = RelationshipFrom('SecurityGroup', 'APPLIED')
15+
OwnerId = StringProperty()
16+
VpcId = StringProperty()
17+
18+
@staticmethod
19+
def arn(region, owner_id, cluster_id):
20+
return 'arn:aws:redshift:{}:{}:cluster/{}'.format(region, owner_id, cluster_id)
21+
22+
def load():
23+
redshift = boto3.client('redshift')
24+
OwnerId, AWSRegion = get_account_id_and_region()
25+
paginator = redshift.get_paginator('describe_clusters')
26+
for page in paginator.paginate():
27+
for db in page['Clusters']:
28+
name = db['ClusterIdentifier']
29+
DBName = db['DBName']
30+
arn = RedshiftCluster.arn(AWSRegion, OwnerId, name)
31+
properties = { 'DBName': DBName, 'ARN': arn, 'OwnerId': OwnerId,
32+
'VpcId': db['VpcId'], 'ClusterInstanceIdentifier': name}
33+
cluster = RedshiftCluster.create_or_update(properties)[0]
34+
sys.stderr.write('INFO: loading redshift cluster {}\n'.format(name))
35+
36+
# remove old associations
37+
for group in cluster.SecurityGroups:
38+
cluster.SecurityGroups.disconnect(group)
39+
40+
for GroupId in map(lambda g: g['VpcSecurityGroupId'], db['VpcSecurityGroups']):
41+
group = SecurityGroup.get_or_create({'ARN': SecurityGroup.arn(AWSRegion, OwnerId, GroupId)})[0]
42+
sys.stderr.write('INFO: applying security group {} to cluster {}\n'.format(group.GroupId, name))
43+
cluster.SecurityGroups.connect(group)
44+

‎aws_visualizer/neo4j/security_group.py

+80
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
import os
2+
import boto3
3+
from neomodel import *
4+
import logging
5+
6+
7+
8+
class Permission(StructuredRel):
9+
FromPort = IntegerProperty()
10+
ToPort = IntegerProperty()
11+
IpProtocol = StringProperty()
12+
13+
class CIDR(StructuredNode):
14+
CidrIp = StringProperty(unique=True)
15+
grants = RelationshipFrom('SecurityGroup', 'GRANTED_BY', model=Permission)
16+
17+
class SecurityGroup(StructuredNode):
18+
ARN = StringProperty(unique_index=True)
19+
OwnerId = StringProperty()
20+
GroupId = StringProperty()
21+
GroupName = StringProperty()
22+
VpcIp = StringProperty()
23+
CidrIpPermissions = RelationshipTo('CIDR', 'GRANTED_TO', model=Permission)
24+
SecurityGroupPermissions = RelationshipTo('SecurityGroup', 'GRANTED_TO', model=Permission)
25+
26+
@staticmethod
27+
def arn(region, owner_id, group_id):
28+
return 'arn:aws:ec2:{}:{}:security-group/{}'.format(region, owner_id, group_id)
29+
30+
31+
def load():
32+
log = logging.getLogger()
33+
groups = {}
34+
ec2 = boto3.client('ec2')
35+
paginator = ec2.get_paginator('describe_security_groups')
36+
pages = paginator.paginate()
37+
for page in pages:
38+
for sg in page['SecurityGroups']:
39+
arn = SecurityGroup.arn(ec2.meta.region_name, sg['OwnerId'], sg['GroupId'])
40+
groups[arn] = sg
41+
42+
for arn, sg in groups.items():
43+
VpcId = sg['VpcId'] if 'VpcId' in sg else None
44+
properties = {n: sg[n] for n in filter(lambda n: n in sg, ['GroupName', 'OwnerId', 'GroupId', 'VpcId'])}
45+
properties['ARN'] = arn
46+
47+
security_group = SecurityGroup.create_or_update(properties)[0]
48+
sys.stderr.write('storing security group {}\n'.format(security_group.GroupId))
49+
50+
for p in security_group.CidrIpPermissions:
51+
security_group.CidrIpPermissions.disconnect(p)
52+
for p in security_group.SecurityGroupPermissions:
53+
security_group.SecurityGroupPermissions.disconnect(p)
54+
55+
56+
for arn, sg in groups.items():
57+
security_group = SecurityGroup.get_or_create({'ARN': arn})[0]
58+
for p in sg['IpPermissions']:
59+
permission = {n : p[n] for n in filter(lambda n : n in p, ['IpProtocol', 'FromPort', 'ToPort'])}
60+
if VpcId is not None:
61+
permission['VpcId'] = VpcId
62+
63+
for c in p['IpRanges']:
64+
CidrIp = CIDR.get_or_create({'CidrIp': c['CidrIp']})[0]
65+
security_group.CidrIpPermissions.connect(CidrIp, permission)
66+
67+
for c in p['Ipv6Ranges']:
68+
CidrIp = CIDR.get_or_create({'CidrIp': c['CidrIpv6']})[0]
69+
security_group.CidrIpPermissions.connect(CidrIp, permission)
70+
71+
for g in p['UserIdGroupPairs']:
72+
properties = {}
73+
properties['ARN'] = SecurityGroup.arn(ec2.meta.region_name, g['UserId'], g['GroupId'])
74+
properties['OwnerId'] = g['UserId']
75+
properties['GroupId'] = g['GroupId']
76+
source_sg = SecurityGroup.get_or_create(properties)[0]
77+
sys.stderr.write('grant IP permission {} from {} to {}\n'.format(permission, security_group.GroupId, source_sg.GroupId))
78+
security_group.SecurityGroupPermissions.connect(source_sg, permission)
79+
80+

‎aws_visualizer/neo4j/sts.py

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
import boto3
2+
3+
def get_account_id_and_region():
4+
sts = boto3.client('sts')
5+
return sts.get_caller_identity()['Account'], sts.meta.region_name
6+

0 commit comments

Comments
 (0)