traffic_

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# This plugin monitors average traffic for one or more ports. It supports IPv4,
# IPv6, TCP and UDP. The plugin assumes iptables and ip6tables chains identical
# to those created by our fw-rules script[1]. See below on how to create those
# chains manually.
#
#
# CONFIGURATION
#
# This example configuration monitors port 80 and 443 on tcp and both IPv4 and
# IPv6:
#     [traffic_http]
#    env.ports 80 443
#    env.title http(s)
#    env.udp n
#    env.tcp y
#    env.ipv4 y
#    env.ipv6 y
#
# 'y' is always the default, so you can just exclude what you don't like. This
# is identical to the above example:
#     [traffic_http]
#    env.ports 80 443
#    env.title http(s)
#    env.udp n
#
# Note again that the appropriate firewall rules *must* be present, otherwise
# they are silently ignored. If, with the above example, port 80 is not counted
# in the appropriate chains, it will not show up in the output.
#
#
# USE AS NON-WILDCARD PLUGIN
#
# If you use this plugin as a normal non-wildcard plugin (that is, name the
# symlink just 'traffic'), it will print the counters for the count-chains
# as a whole.
#
#
# CREATING IPTABLES RULES MANUALLY
#
# If you don't want to use fw-rules[1], you can always create the appropriate
# chains manually. The chains follow the naming scheme:
#    COUNT_<proto>_<dir>
#
# so if you want to count TCP packets, add these firewall rules to the bottom
# of your firewall script:
#    iptables --new-chain COUNT_TCP_IN
#    iptables -I INPUT -p tcp -j COUNT_TCP_IN -i eth0
#    iptables --new-chain COUNT_TCP_OUT
#    iptables -I OUTPUT -p tcp -j COUNT_TCP_OUT -o eth0
#
# Then you must add a rule for each port you want to monitor:
#    iptables -A COUNT_TCP_IN -p tcp --dport 80
#    iptables -A COUNT_TCP_OUT -p tcp --sport 80
#
# The same rules also work with ip6tables.
#
# [1] http://git.fsinf.at/fsinf/fw-rules/
#
#%# family=manual
#%# capabilities=

import sys
import os

if len(sys.argv) > 1 and sys.argv[1] == 'autoconf':
    print('no')
    sys.exit(0)


def get_bool(val):
    if val.lower() in ['n', 'no', 'false']:
        return False
    else:
        return True


class chain():
    def __init__(self, line):
        info = line.split()
        self.name = info[1]
        self.rules = []

    def add_header(self, header):
        self.names = header.split()

    def add_rule(self, raw_rule):
        fields = raw_rule.split(None, len(self.names) - 1)
        self.rules.append(dict(zip(self.names, fields)))

    def get_bytes(self, destination, ip_version, protocol):
        """
        Get bytes on a specific destination. Due to disambiguities in
        iptables output, we have to prepend the protocol with ipversion
        'ipv4'.
        """
        if ip_version == 'ipv4':
            destination = '%s %s' % (protocol, destination)

        for rule in self.rules:
            if rule['destination'] == destination:
                return int(rule['bytes'])
        return None

    def get_bytes_new(self, ip_version, target=None, destination=None):
        for rule in self.rules:
            if target and rule['target'] != target:
                continue
            if destination and rule['destination'] != destination:
                continue
            return int(rule['bytes'])
        return -1

# see what we want to count:
use_v6 = get_bool(os.environ.pop('ipv6', 'y'))
use_v4 = get_bool(os.environ.pop('ipv4', 'y'))
use_tcp = get_bool(os.environ.pop('tcp', 'y'))
use_udp = get_bool(os.environ.pop('udp', 'y'))
maximum = int(os.environ.get('max', '1073741824'))
if '_' not in sys.argv[0]:
    ports = []
else:
    ports = os.environ.pop('ports', '').split()
    if not ports:
        print('Error: You need to set env.ports')
        sys.exit(1)

from subprocess import Popen, PIPE


def parse_tables(binary):
    p = Popen([binary, '-L', '-w', '-n', '-v', '-x'], text=True, stdout=PIPE)
    stdout = p.communicate()[0].splitlines()
    if p.returncode != 0:
        return {}

    chains = {}
    cur_chain = None
    for line in stdout:
        l = line.lower().strip()
        if not l:  # empty line
            continue
        elif l.startswith('chain'):  # new chain definition
            cur_chain = chain(l)
            chains[cur_chain.name] = cur_chain
            continue
        elif l.startswith('pkts'):  # header line
            cur_chain.add_header(l)
        else:
            cur_chain.add_rule(l)
    return chains

v6_total = 0
v4_total = 0
in_total = 0
out_total = 0


def print_traffic(ports, protocol, ip_version, chains):
    printed_in = 0
    printed_out = 0

    if ports:
        in_chain = chains.pop('count_%s_in' % protocol, None)
        out_chain = chains.pop('count_%s_out' % protocol, None)

        for port in ports:
            prefix = '%s_%s_%s_' % (protocol, port, ip_version)
            if in_chain:
                in_bytes = in_chain.get_bytes(
                    'dpt:%s' % port, ip_version, protocol)
                if in_bytes != None:
                    print('%sin.value %s' % (prefix, in_bytes))
                    printed_in += in_bytes
            if out_chain:
                out_bytes = out_chain.get_bytes(
                    'spt:%s' % port, ip_version, protocol)
                if out_bytes != None:
                    print('%sout.value %s' % (prefix, out_bytes))
                    printed_out += out_bytes

        return printed_in, printed_out
    else:
        in_chain = chains['input']
        out_chain = chains['output']

        prefix = '%s_%s' % (protocol, ip_version)
        target_in = 'count_%s_in' % protocol
        target_out = 'count_%s_out' % protocol

        in_bytes = in_chain.get_bytes_new(ip_version, target=target_in)
        out_bytes = out_chain.get_bytes_new(ip_version, target=target_out)

        print('%s_in.value %s' % (prefix, in_bytes))
        print('%s_out.value %s' % (prefix, out_bytes))

        return in_bytes, out_bytes


def config(field, label):
    print('''%s.label %s
%s.type DERIVE
%s.min 0''' % (field, label, field, field))
    if maximum:
        print('%s.max %s' % (field, maximum))


def print_traffic_config(ports, protocol, ip_version, chains):
    in_chain = chains.pop('count_%s_in' % protocol, None)
    out_chain = chains.pop('count_%s_out' % protocol, None)
    if ip_version == 'ipv4':
        info_proto = 'IPv4'
    elif ip_version == 'ipv6':
        info_proto = 'IPv6'

    if ports:
        for port in ports:
            prefix = '%s_%s_%s_' % (protocol, port, ip_version)
            if in_chain:
                in_bytes = in_chain.get_bytes(
                    'dpt:%s' % port, ip_version, protocol)
                if in_bytes != None:
                    field = '%sin' % prefix
                    label = 'Port %s/%s (%s, incoming)' % (port, protocol, info_proto)
                    config(field, label)
            if out_chain:
                out_bytes = out_chain.get_bytes(
                    'spt:%s' % port, ip_version, protocol)
                if out_bytes != None:
                    field = '%sout' % prefix
                    label = 'Port %s/%s (%s, outgoing)' % (port, protocol, info_proto)
                    config(field, label)
    else:
        if use_tcp and use_v4:
            config('tcp_ipv4_in', 'TCP/IPv4 incoming')
            config('tcp_ipv4_out', 'TCP/IPv4 outgoing')
        if use_tcp and use_v6:
            config('tcp_ipv6_in', 'TCP/IPv6 incoming')
            config('tcp_ipv6_out', 'TCP/IPv6 outgoing')
        if use_udp and use_v4:
            config('udp_ipv4_in', 'UDP/IPv4 incoming')
            config('udp_ipv4_out', 'UDP/IPv4 outgoing')
        if use_udp and use_v6:
            config('udp_ipv6_in', 'UDP/IPv6 incoming')
            config('udp_ipv6_out', 'UDP/IPv6 outgoing')


if len(sys.argv) > 1 and sys.argv[1] == 'config':
    title = os.environ.pop('title', sys.argv[0].partition('_')[2])
    if title:
        title = ' via %s' % title
    print("""graph_title Traffic%s
graph_vlabel bytes/s
graph_category network
graph_args -l 0""" % title)
    if use_v6:
        v6_chains = parse_tables('ip6tables')
        if use_tcp:
            print_traffic_config(ports, 'tcp', 'ipv6', v6_chains)
        if use_udp:
            print_traffic_config(ports, 'udp', 'ipv6', v6_chains)
    if use_v4:
        v4_chains = parse_tables('iptables')
        if use_tcp:
            print_traffic_config(ports, 'tcp', 'ipv4', v4_chains)
        if use_udp:
            print_traffic_config(ports, 'udp', 'ipv4', v4_chains)

    if use_v4 and use_v6:
        config('total_ipv4', 'Total traffic via IPv4')
        config('total_ipv6', 'Total traffic via IPv6')
    if len(ports) > 1:
        config('total_in', 'Total incoming traffic')
        config('total_out', 'Total outgoing traffic')

    config('total', 'Total traffic')
    sys.exit(0)

if use_v6:
    v6_chains = parse_tables('ip6tables')
    if use_tcp:
        i, o = print_traffic(ports, 'tcp', 'ipv6', v6_chains)
        v6_total += i + o
        in_total += i
        out_total += o
    if use_udp:
        i, o = print_traffic(ports, 'udp', 'ipv6', v6_chains)
        v6_total += i + o
        in_total += i
        out_total += o
if use_v4:
    v4_chains = parse_tables('iptables')
    if use_tcp:
        i, o = print_traffic(ports, 'tcp', 'ipv4', v4_chains)
        v4_total += i + o
        in_total += i
        out_total += o
    if use_udp:
        i, o = print_traffic(ports, 'udp', 'ipv4', v4_chains)
        v4_total += i + o
        in_total += i
        out_total += o

if use_v4 and use_v6:
    print('total_ipv4.value %s' % v4_total)
    print('total_ipv6.value %s' % v6_total)

if len(ports) > 1:
    print('total_in.value %s' % in_total)
    print('total_out.value %s' % out_total)

total = v4_total + v6_total
print('total.value %s' % total)