kirikae/python-nettopo

Network topology from netstat output

python-nettopo's Issues

Updates suggested

#!/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()