#!/usr/bin/env python
import socket
import argparse
from getpass import getpass
from collections import namedtuple
from neo4j import GraphDatabase
##########################
# Create namedtuples used
# ------------------------
Connection = namedtuple('Connection', ['proto', 'recvq', 'sendq', 'localaddr', 'localport', 'remoteaddr', 'remoteport', 'state', 'pid', 'program'])
Host = namedtuple('Host', ['hostname', 'ipaddr'])
##########################
# Constants
# ------------------------
URI = "bolt://localhost:7687"
##########################
# Parse Arguments
# ------------------------
parser = argparse.ArgumentParser(
description = "Parse Netstat output files and import them into Neo4j Graph Database"
)
parser.add_argument('--files', '-f', nargs='+', help='Pass the File or Files to analyse')
args = parser.parse_args()
for arg in vars(args):
filenames = getattr(args, arg)
# Create driver for neo4j:
# TODO: Allow these to be parsed in from STDIO or ENV
neo4jusername = input("Neo4j username: ")
neo4jpassword = getpass("Neo4j password: ")
driver = GraphDatabase.driver(URI, auth=(neo4jusername, neo4jpassword))
class CachedDNSLookup:
"""
This class provides a simple dictionary backed cache for hostname lookups.
"""
def init(self):
self._cache= {}
def get_host_by_address(self, addr):
if addr not in self._cache.keys():
self._cache[addr] = socket.gethostbyaddr(addr)[0]
return self._cache[addr]
def main():
"""
The main entry point!
:return: None
"""
all_connections = []
sources = {}
dns_lookup = CachedDNSLookup()
# First, we read this into an array, parsing as we go.
for file in filenames:
print("Reading in file: ",file)
with open(file, 'r') as f:
for line in f:
current_line = line.split()
print(current_line)
if current_line[0] == "tcp":
(proto, recvq, sendq, local, remote, state, pidprog) = current_line
elif current_line[0] == "udp":
if current_line[5] == "ESTABLISHED":
(proto, recvq, sendq, local, remote, state, pidprog) = current_line
else:
state = ""
(proto, recvq, sendq, local, remote, pidprog) = current_line
else:
continue
if pidprog == '-':
(pid, program) = ("UNKNOWN", "UNKNOWN")
else:
(pid, program) = pidprog.split('/')
(localaddr, localport) = local.split(':')[:2]
(remoteaddr, remoteport) = remote.split(':')[:2]
current_connection = Connection(proto, recvq, sendq, localaddr, localport, remoteaddr, remoteport, state, pid, program)
all_connections.append(current_connection)
print("Finished with :", file)
# Create working Sets
all_addresses = set()
local_connection = set()
remote_connections = set()
all_programs = set()
for conn in all_connections:
all_addresses.add(conn.localaddr)
all_addresses.add(conn.remoteaddr)
local_connection.add(conn.localaddr)
remote_connections.add(conn.remoteaddr)
all_programs.add(conn.program)
print("Connections loaded.")
with driver.session() as session:
for host in all_addresses:
try:
name = dns_lookup.get_host_by_address(host)
except:
name = "UNKNOWN"
add_host_str = 'CREATE (A:COMPUTER {IP: "{}", FQDN: "{}"})'.format(host, name)
session.run(add_host_str)
print("Hosts loaded.")
with driver.session() as session:
for program in all_programs:
add_program = 'CREATE (A:PROGRAM {Name: "{}"})'.format(program)
session.run(add_program)
print("Programs loaded.")
with driver.session() as session:
ADD_APP_STR = 'MATCH (A:COMPUTER {IP: "{}"}),(B:PROGRAM {Name: "{}"}),(C:COMPUTER {IP: "{}"}) CREATE (A)-[:RUNS]->(B)'
ADD_HOST_REL_STR = 'MATCH (A:COMPUTER {IP: "{}"}),(B:PROGRAM {Name: "{}"}),(C:COMPUTER {IP: "{}"}) CREATE (B)-[:CONNECTS_TO {Local_Port: "{}", Remote_Port: "{}", Protocol: "{}", State: "{}"}]->(C)'
for conn in all_connections:
if conn.localaddr not in sources.keys():
sources[conn.localaddr] = []
if conn.remoteaddr not in sources[conn.localaddr]:
add_app_relationship = ADD_APP_STR.format(conn.localaddr, conn.program, conn.remoteaddr)
add_host_relationship = ADD_HOST_REL_STR.format(conn.localaddr, conn.program, conn.remoteaddr, conn.localport, conn.remoteport, conn.proto, conn.state)
session.run(add_app_relationship)
session.run(add_host_relationship)
sources[conn.localaddr].append(conn.remoteaddr)
print(add_app_relationship)
print(add_host_relationship)
print("Connections loaded.")
if name == 'main':
main()