#!/usr/bin/python3

import os
import sys
import psutil
import re
import shutil
from socket import AF_INET, AF_INET6, SOCK_STREAM, SOCK_DGRAM
from pprint import pprint

# read parameters
prefix = os.getenv('PREFIX')
status = os.getenv('STATUS', 'LISTEN')
laddr = os.getenv('LADDR', '127.0.0.1')
skip_exe = os.getenv('SKIP_EXE','')
skip_port = list(map(int, filter(None, os.getenv('SKIP_PORT','').split(' '))))
format = os.getenv('FORMAT','{port}')
dformat = os.getenv('DFORMAT', '{port}:{name}')

fails = 0

proto_map = {
    (AF_INET, SOCK_STREAM): 'TCP',
    (AF_INET6, SOCK_STREAM): 'TCP6',
    (AF_INET, SOCK_DGRAM): 'UDP',
    (AF_INET6, SOCK_DGRAM): 'UDP6',
}

outdata=list()
ports = list()

for p in psutil.process_iter():
    try:
        conns = p.connections()
        if len(conns):

            for c in conns:
                if not re.search(status, c.status):
                    continue

                if not re.search(laddr, c.laddr[0]):
                    continue

                if skip_exe and re.search(skip_exe, p.exe()):
                    continue

                if c.laddr[1] in skip_port:
                    continue

                if c.laddr[1] in ports:
                    continue

                ports.append(c.laddr[1])

                cs = {}
                # cs['iname']=ts.iname
                cs['name'] = p.name()
                cs['pid'] = p.pid
                cs['basename'] = os.path.basename(p.exe())
                cs['exe'] = p.exe()
                cs['proto'] = proto_map[(c.family, c.type)]
                cs['ip'] = c.laddr[0]
                cs['port'] = c.laddr[1]
                cs['status'] = c.status

                outdata.append(cs)

    except (psutil.Error, psutil.AccessDenied, psutil.NoSuchProcess):
        fails += 1
        pass

portstr = ' '.join(sorted((format.format(**cs) for cs in outdata), key=int))
d = ' '.join(sorted(dformat.format(**cs) for cs in outdata))
print("NAME: {}opentcp".format(prefix))
print("TAGS: opentcp")
print("METHOD: string|options=reinit dynamic")
print("DETAILS: {} ({} fails)".format(d, fails))

print("STATUS: {}".format(portstr))
