In [ ]:
import datetime
import struct, socket
import numpy as np
import linecache, bisect
import csv
import operator
import json
import os
import pandas as pd

try:
    import ipywidgets as widgets # For jupyter/ipython >= 1.4
except ImportError:
    from IPython.html import widgets
from IPython.display import display, Javascript, clear_output

spath = os.getcwd()
path = spath.split("/") 
date = path[len(path)-1]   
dpath = '/'.join(['data' if var == 'ipynb' else var for var in path]) + '/'
cpath = '/'.join(['context' if var == 'ipynb' else var for var in path][:len(path)-2]) + '/'

anchor = ''
ir_f = ''
threat_name = ''
iplist = ''
top_results = 20
ip_comments = {}

In [ ]:
# Widget styles and initialization
topBox = widgets.Box()
bottomBox = widgets.Box()
mainBoxes_css = (
    (None, 'width', '90%'),
    (None, 'margin', '0 auto'),
)

topBox._css = mainBoxes_css
bottomBox._css = mainBoxes_css 

separator = widgets.HBox(width='100%', height='20px')
threatBox = widgets.HBox(width='100%', height='auto')
threat_title = widgets.HTML(height='25px', width='100%')
threat_list_container = widgets.Box(width='80%', height='100%')
threat_button_container = widgets.Box(width='20%', height='100%')
susp_select = widgets.Select(height='100%', width='99%')
search_btn = widgets.Button(description='Search',height='100%', width='65px')
search_btn.button_style = 'primary'
susp_select._css = (
    (None, 'height', '90%'),
    (None, 'width', '95%'),
    ('select', 'overflow-x', 'auto'),
    ('select', 'margin', 0)
)

resultSummaryBox = widgets.Box()
result_title = widgets.HTML(width='100%')
result_summary_box = widgets.HBox(width='100%')
result_summary_container = widgets.Box(width='80%')
result_button_container =  widgets.Box(width='20%')
result_summary_box.children = [result_title, result_summary_container, result_button_container]
 
resultTableBox = widgets.HBox()
result_box_css = (
    (None, 'overflow', 'hidden'),
    (None, 'width', '100%'),
)

resultSummaryBox._css = result_box_css
resultTableBox._css = result_box_css
 
result_content_box_css = (
    (None, 'overflow','auto'), 
    (None, 'width', '50%'),
)

threat_button_container._css = (
    (None, 'padding-top', '30px'), 
)  

topBox.children = [threatBox]
bottomBox.children = [resultSummaryBox,resultTableBox]

Functions Definition


In [ ]:
top_inbound_b='' 
top_outbound_b=''
top_twoway_b=''

inbound='' 
outbound=''
twoway=''

global src_per_conns
global src_per_bytes
global dst_per_conns
global dst_per_bytes

def start_investigation(): 
    display(Javascript("$('.widget-area > .widget-subarea > *').remove();"))   
    external_ips = []
    clear_output()
    global ip_comments
    
    response = GraphQLClient.request(
        query="""query($date:SpotDateType!) {
                flow{
                    threats{
                        list(date:$date) {
                           srcIp
                           dstPort
                           dstIp
                           srcPort
                           score 
                        }
                    }
            }
        }""",
        variables={
            'date': datetime.datetime.strptime(date, '%Y%m%d').strftime('%Y-%m-%d')
        }
    ) 
     
    ip_comments = GraphQLClient.request(
        query="""query($date:SpotDateType!) {
                flow{
                    threats{
                        comments(date:$date) {
                            ip
                            title
                            text
                        }
                    }
            }
        }""",
        variables={
            'date': datetime.datetime.strptime(date, '%Y%m%d').strftime('%Y-%m-%d')
        }
    )
    ip_comments = ip_comments['data']['flow']['threats']['comments']
    if not 'errors' in response : 
        for row in response['data']['flow']['threats']['list']:
            if row['score'] == 1: 
                srcIP = ''
                dstIP = '' 
                if row['srcIp'] not in external_ips: 
                    external_ips.append(row['srcIp'])
                if row['dstIp'] not in external_ips: 
                    external_ips.append(row['dstIp'])

        if len(external_ips) == 0:
            display(widgets.Box((widgets.HTML(value="There are no connections scored as High risk.\
                You can score some connections at the 'Suspicious' panel.", width='100%'),)))
        else:  
            sorted_dict = sorted(external_ips, key=operator.itemgetter(0))      
            display_controls(sorted_dict)  
    else:   
        
        display(widgets.Box((widgets.HTML(value="An error occurred while trying to get the results:" 
                                          + response['errors'][0]['message'], width='100%'),)))
        

def display_controls(threat_list):
    threat_title.value ="<h4>Suspicious Connections</h4>"    
    susp_select.options = threat_list
    susp_select.height=150
    susp_select.selected_label = threat_list[0]
    threat_list_container.children = [threat_title,susp_select]
    threat_button_container.children = [search_btn]
    threatBox.children = [threat_list_container, threat_button_container]
    display(topBox) 
    
    def search_ip(b):
        global anchor 
        global top_inbound_b
        global expanded_results
        anchor = susp_select.value   
        if anchor != "":
            clear_output()        
            removeWidget(1)
            
            expanded_results = GraphQLClient.request(
                        query="""query($date:SpotDateType!,$ip:SpotIpType!){
                                  flow{
                                    threat{
                                      details(date:$date,ip:$ip) {
                                        srcIp
                                        maxBytes
                                        connections
                                        maxPkts
                                        avgPkts
                                        lastSeen
                                        srcPort
                                        firstSeen 
                                        dstIp
                                        avgBytes
                                        dstPort
                                      }
                                    }
                                  }  
                                }
                                """,
                            variables={
                            'date': datetime.datetime.strptime(date, '%Y%m%d').strftime('%Y-%m-%d'),
                            'ip': anchor
                            }
                        ) 
             
            if not 'errors' in expanded_results :  
                print "\n Looking for additional details..."
                display_threat_box(anchor)

                get_in_out_and_twoway_conns() 
                 
                display(bottomBox)     
            else: 
                display(widgets.Box((widgets.HTML(value="Something went wrong. \
                    The expanded search couldn't be performed" + expanded_results['errors'][0]['message'], width='100%'),)))
                        
    search_btn.on_click(search_ip)


def display_threat_box(ip):
    clear_output()
    global ip_comments    
    title ="" 
    text = ""
    title = next((item['title'] for item in ip_comments if item.get("ip") == ip), "")
    text = next((item['text'] for item in ip_comments if item.get("ip") == ip), "")  
    result_title.value="<h4 class='spot-text-wrapper spot-text-xlg' data-toggle='tooltip'>Threat summary for " + anchor +"</h4>"
    tc_txt_title = widgets.Text(value=title, placeholder='Threat Title', width='100%')
    tc_txa_summary = widgets.Textarea(value=text, height=100, width='95%')
    tc_btn_save = widgets.Button(description='Save', width='65px', layout='width:100%')
    tc_btn_save.button_style = 'primary'
     
    tc_txt_title._css = (
        (None, 'width', '95%'),
    )
    
    result_summary_container.children = [tc_txt_title, tc_txa_summary]
    result_button_container.children=[tc_btn_save]
    result_summary_box.children = [result_summary_container, result_button_container]
    resultSummaryBox.children = [result_title,result_summary_box]
    
    def save_threat_summary(b):
        clear_output()        
        removeWidget(1) 
        response = ""
        
        response += add_threat(anchor, tc_txt_title.value, tc_txa_summary.value.replace('\n', '\\n'))
        response += "Story board successfully created for {0}".format(anchor)
        start_investigation()
        display(widgets.Box((widgets.HTML(value=response, width='100%'),)))
        
        
    tc_btn_save.on_click(save_threat_summary)

In [ ]:
def get_in_out_and_twoway_conns():
    srcdict = {}
    dstdict = {}
    conns_dict= {} 
    rowct = 0 
    if not 'errors' in expanded_results :         
        df = pd.DataFrame(expanded_results['data']['flow']['threat']['details']) 
        for row in expanded_results['data']['flow']['threat']['details']:  
            srcdict[row['srcIp']] = {
                'ip_int': struct.unpack("!L", socket.inet_aton(str(row['srcIp'])))[0],
                'dst_ip': row['dstIp'],
                'dst_ip_int': struct.unpack("!L", socket.inet_aton(str(row['dstIp'])))[0],
                'conns': int(row['connections']),
                'maxbytes': int(row['maxBytes'])
            }
            dstdict[row['dstIp']] = {
                 'ip_int': struct.unpack("!L", socket.inet_aton(str(row['dstIp'])))[0],
                 'src_ip': row['srcIp'],
                 'src_ip_int': struct.unpack("!L", socket.inet_aton(str(row['srcIp'])))[0],
                 'conns': int(row['connections']),
                 'maxbytes': int(row['maxBytes'])
                }       
            rowct +=1  
        
        src = df.loc[df['dstIp'] == anchor]
        src_per_conns = src.sort_values('connections',0,False)
        src_per_bytes = src.sort_values('maxBytes',0,False)
        dst = df.loc[df['srcIp'] == anchor]
        dst_per_conns = dst.sort_values('connections',0,False)
        dst_per_bytes = dst.sort_values('maxBytes',0,False)

        children = []
        children += (display_results(['srcIp','connections','srcPort','dstPort'], src_per_conns, 
                                          top_results),)
        children += (display_results(['dstIp','connections','srcPort','dstPort'], dst_per_conns, 
                                          top_results),)
        children += (display_results(['srcIp','maxBytes','srcPort','dstPort'], src_per_bytes, 
                                           top_results),)
        children += (display_results(['dstIp','maxBytes','srcPort','dstPort'], dst_per_bytes, 
                                           top_results),)

        result_tabs = widgets.Accordion(children=children, width='100%', selected_index=-1)
        result_tabs.set_title(0,"Top source IP per connections")
        result_tabs.set_title(1,"Top destination IP per connections")
        result_tabs.set_title(2,"Top source IP per bytes transferred")
        result_tabs.set_title(3,"Top destination IP per bytes transferred")
        
        result_tabs._css = (
            (None, 'margin-top', '10px'),
            (None, 'margin-bottom', '10px'),
        )

        resultTableBox.children = [result_tabs,]
         

def display_results(cols, dataframe, top):    
    table = dataframe[:top].to_html(classes='table table-striped table-bordered table-hover', columns=cols, index=False)
    return widgets.HTML(value=table, width='100%')


    
def add_threat(ip,threat_title, threat_comment):
       
    mutation="""mutation(
                $date: SpotDateType, 
                $ip: SpotIpType!, 
                $text: String!, 
                $title: String!,
                $threatDetails: [NetflowThreatDetailsInputType]!,
                $topResults:Int) 
                {
                  flow {
                    createStoryboard(input:{
                        threatDetails: $threatDetails,
                        date: $date, 
                        ip: $ip, 
                        title: $title, 
                        text: $text,
                        first:$topResults})
                    {success}
                  }
                }"""
 
    variables={
        'date': datetime.datetime.strptime(date, '%Y%m%d').strftime('%Y-%m-%d'),
        'ip': ip,
        'title': threat_title,
        'text': threat_comment,
        'threatDetails': expanded_results['data']['flow']['threat']['details'],
        'first':top_results

        }
    
    response = GraphQLClient.request(mutation,variables)
    if not 'errors' in response:
        return "Story board successfully created"
    else:
        return response['errors'][0]['message']
    
def removeWidget(index):
    js_command = "$('.widget-area > .widget-subarea > .widget-box:eq({0})').remove();".format(index)    
    display(Javascript(js_command))

In [ ]:
start_investigation()