!pip install graphistry -q
!pip install awscli -q

!aws configure set region us-west-2
!aws configure set aws_access_key_id "FILL_ME_IN"
!aws configure set aws_secret_access_key "FILL_ME_IN"

import pandas as pd
import json
import graphistry

# To specify Graphistry account & server, use:
# graphistry.register(api=3, username='...', password='...', protocol='https', server='hub.graphistry.com')
# For more options, see https://github.com/graphistry/pygraphistry#configure


"""
Explanation: AWS CloudWatch VPC Flow Logs <> Graphistry
Analyze cloudwatch logs, such as using vpc flow to map an account, with Graphistry
This example directly uses the AWS CLI for cloudwatch API access. You can also work from S3 or systems like Athena.
Installs & Configure
Set aws_access_key_id, aws_secret_access_key, key or pull from your env
End of explanation
"""


!aws logs describe-log-groups

!aws logs filter-log-events --log-group-name VPCFlowDemo > data.json
!ls -al data.json

with open('data.json', 'r') as f:
    data = json.load(f)
df = pd.DataFrame([x['message'].split(" ") for x in data['events']])
df.columns = cols = ['version', 'accountid', 'interfaceid', 'src_ip', 'dest_ip', 'src_port', 'dest_port', 'protocol', 'packets', 'bytes', 'time_start', 'time_end', 'action', 'status']

print('# rows', len(df))
df.sample(3)

# Int->Float for precision errors
df2 = df.copy()
for c in ['packets', 'bytes']:
    df2[c] = df2[c].astype(float)

summary_df = df2\
    .groupby(['src_ip', 'dest_ip', 'interfaceid', 'dest_port', 'protocol', 'action', 'status'])\
    .agg({
        'time_start': ['min', 'max'],
        'time_end': ['min', 'max'],
        'packets': ['min', 'max', 'sum', 'count'],
        'bytes': ['min', 'max', 'sum', 'count']
    }).reset_index()
summary_df.columns = [(" ".join(x)).strip().replace(" ", "_") for x in list(summary_df.columns)]
print('# rows', len(summary_df))
summary_df.sample(3)

"""
Explanation: Record logs
If you do not already have logs, you can record VPC flow logs from your EC2 console:
  * Services -> EC2 -> Network Interfaces -> select interface(s) ->  Action -> create flow log
    * Send to cloudwatch; use default settings for IAM and elsewhere
  * When enough data available, stop logging
Download & summarize logs

Pick a log group from available
Fetch: See AWS docs on filter-log-events
Load into a dataframe
Compute summary stats
End of explanation
"""


hg = graphistry.hypergraph(
    summary_df,
    entity_types=['src_ip', 'dest_ip'], #'dest_port', 'interfaceid', 'action', ...
    direct=True)
hg['graph'].bind(edge_title='bytes_sum').plot()

"""
Explanation: Plot
End of explanation
"""


