In [ ]:

    

import urllib2
import json
import os
import csv

# getting date from the parent path. 
path = os.getcwd().split("/") 
date = path[len(path)-1]   
dsource = path[len(path)-2]  
dpath = '/'.join(['data' if var == 'ipynb' else var for var in path]) + '/'

sconnect = dpath + 'dns_scores.csv'
sconnectbu = dpath + 'dns_scores_bu.csv'
score_tmp = dpath + 'score_tmp.csv'  
score_fbk = dpath + 'dns_scores_fb.csv'


    

In [ ]:

    

def apply_css_to_select(select):
    select._css = (
        (None, 'height', '90%'),
        (None, 'width', '90%'),
        ('select', 'overflow-x', 'auto'),
        ('select', 'width', '100%'),
        ('select', 'margin', 0)
    )

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

def fill_list(list_control,source):
    options_list = ['--Select--'] 
    options_list.extend([s for s in source])
    list_control.options = options_list

# client panel
client_header = widgets.HTML(value="Client IP")
client_select = widgets.Select(height='90%')
apply_css_to_select(client_select)

client_box = widgets.Box(width='20%', height='100%')
client_box.children = [client_header, client_select]

# query panel
query_header = widgets.HTML(value="Query")
query_select = widgets.Select(height='90%')
apply_css_to_select(query_select)

query_box = widgets.Box(width='60%', height='100%')
query_box.children = [query_header, query_select]

# Actions Panel
actions_header = widgets.HTML(value=" ")
quick_text = widgets.Text(value='', width='100%', placeholder='Quick scoring')
quick_text._css = (
    (None, 'width', '100%'),
)
rating_btn = widgets.RadioButtons(description='Rating:', options=['1', '2', '3'], width='100%')
assign_btn = widgets.Button(description='Score', width='45%')
assign_btn.button_style = 'primary'
save_btn = widgets.Button(description='Save', width='45%')
save_btn.button_style = 'primary'
save_btn._css = (
    (None, 'margin-left', '10%'),
)
actions_box = widgets.Box(width='20%', height='100%')
actions_box.children = [actions_header,quick_text,rating_btn, assign_btn,save_btn]

scoring_form = widgets.HBox(width='90%', height=250)
scoring_form.children = [client_box,query_box,actions_box]


def data_loader(): 
    us_ips = []
    us_dns = []

    with open(sconnect, 'r') as f:
        reader = csv.DictReader(f, delimiter=',')
        for row in reader:           
            if row['ip_dst'] not in us_ips and row['ip_sev'] == '0': 
                us_ips.append(row['ip_dst'])
            if row['dns_qry_name'] not in us_dns and row['dns_sev'] == '0':
                us_dns.append(row['dns_qry_name']) 

    fill_list(client_select,us_ips)
    fill_list(query_select,us_dns)
    client_select.value = "--Select--"
    query_select.value = "--Select--"    


display(Javascript("$('.widget-area > .widget-subarea > *').remove();"))
data_loader()
display(scoring_form)


    

In [ ]:

    

import csv
import datetime
import subprocess 

def assign_score(b):
    score_values = []
    scored_threats = []
    ip_sev = int(rating_btn.selected_label) if not "--Select--" in client_select.value else ""
    dns_sev = int(rating_btn.selected_label) if not "--Select--" in query_select.value else ""    

    if quick_text.value: 
        ip = ""
        dns = quick_text.value
        dns_sev = int(rating_btn.selected_label) 
        # Loop over every element in query_select widget
        score_values = []
        for query in query_select.options:
            if query.endswith(dns):
                # Matching element, create one row
                score_values.append((ip,query,ip_sev,dns_sev))
    else: 
        ip = client_select.value if not "--Select--" in client_select.value else ""
        dns = query_select.value if not "--Select--" in query_select.value else ""
        score_values.append((ip,dns,ip_sev,dns_sev))

    with open(sconnect, 'r') as f:
        reader = csv.DictReader(f, delimiter=',')
        rowct = 0
        with open(score_tmp, 'w') as score:
            wr = csv.DictWriter(score, delimiter=',', quoting=csv.QUOTE_NONE, fieldnames=reader.fieldnames)            
            wr.writeheader()
            for row in reader:   
                for value in score_values: 
                    if row['ip_dst'] == value[0]:  
                        row['ip_sev'] = value[2]       
                        scored_threats.append(row)  
                        rowct += 1                  
                        break
                    if row['dns_qry_name'] == value[1]:  
                        row['dns_sev'] = value[3]                        
                        scored_threats.append(row)  
                        rowct += 1
                        break
                wr.writerow(row)     
                    
        if not os.path.exists(score_fbk):  
            with open(score_fbk, 'w') as feedback:
                wr = csv.DictWriter(feedback, delimiter='\t', quoting=csv.QUOTE_NONE, fieldnames=reader.fieldnames)            
                wr.writeheader()

        with open(score_fbk, 'a') as feedback:
            for row in scored_threats:
                wr = csv.DictWriter(feedback, delimiter='\t', quoting=csv.QUOTE_NONE, fieldnames=reader.fieldnames)            
                wr.writerow(row)

    clear_output()
    print "{0} matching connections scored".format(rowct)
    !mv $score_tmp $sconnect 

    if ip != "--Select--":
        display(Javascript("$(\"option[data-value='" + ip +"']\").remove();"))
    if quick_text.value:
        display(Javascript("$(\"option[data-value$='" + quick_text.value +"']\").remove();"))
    elif dns != "--Select--":
        display(Javascript("$(\"option[data-value='" + dns +"']\").remove();"))

    client_select.value = "--Select--"
    query_select.value = "--Select--"
    quick_text.value = ""


def save(b):   
    clear_output()    
    display(Javascript("$('.widget-area > .widget-subarea > *').remove();"))
    data_loader() 
    display(scoring_form)
    display(Javascript('reloadParentData();'))
    ml_feedback() 
    print "Suspicious connects successfully updated"


assign_btn.on_click(assign_score)
save_btn.on_click(save)
        

def ml_feedback():
    dst_name = os.path.basename(sconnect)
    str_fb="DSOURCE={0} &&\
        FDATE={1} &&\
        source /etc/spot.conf &&\
        usr=$(echo $LUSER | cut -f3 -d'/') &&\
        mlnode=$MLNODE &&\
        lpath=$LPATH &&\
        scp {2} $usr@$mlnode:$lpath/{3}".format(dsource,date,score_fbk,dst_name) 

    subprocess.call(str_fb, shell=True)


    

In [ ]:

    

# !cp $sconnectbu $sconnect