import routeros_api
from prometheus_client import start_http_server, Gauge
import time
import re
import os

APP_NAME = 'mktxp-extension'
RTR_ADDRESS = os.environ.get('RTR_ADDRESS')
RTR_USER = os.environ.get('RTR_USER')
RTR_PASS = os.environ.get('RTR_PASS')

gauge_wg_rx = Gauge(APP_NAME+'_wireguard_rx', 'Rx amount on peer', ['peer','interface'])
gauge_wg_tx = Gauge(APP_NAME+'_wireguard_tx', 'Tx amount on peer', ['peer','interface'])
gauge_wg_handshake = Gauge(APP_NAME+'_wireguard_handshake', 'Last peer handshake', ['peer','interface'])
gauge_wg_active = Gauge(APP_NAME+'_wireguard_active', 'If peer is active, is set to inactive after 5 minutes of last connection', ['peer','interface'])
gauge_mangle_rx = Gauge(APP_NAME+'_manle_rx', 'Rx amount on rule', ['comment'])
gauge_mangle_tx = Gauge(APP_NAME+'_mangle_tx', 'Tx amount on rule', ['comment'])

connection = routeros_api.RouterOsApiPool(RTR_ADDRESS, username=RTR_USER, password=RTR_PASS, plaintext_login=True)
api = connection.get_api()

def get_wg():

    list = api.get_resource('/interface/wireguard/peers')
    peers = list.get()
    for peer in peers:
        seconds = 0

        gauge_wg_rx.labels(peer=peer['name'], interface=peer['interface']).set(int(peer['rx']))
        gauge_wg_tx.labels(peer=peer['name'], interface=peer['interface']).set(int(peer['tx']))
        lasthandshake = ''
        try:
            lasthandshake = peer['last-handshake']

            tm = re.split(r'[dhms]',lasthandshake)
            tm_rev = []
            multiply = []
            for i in range(0,len(tm)-1):
                tm_rev.append(tm[len(tm)-2-i])
            if 's' in lasthandshake:
                multiply.append(1)
            if 'm' in lasthandshake:
                multiply.append(60)
            if 'h' in lasthandshake:
                multiply.append(3600)
            if 'd' in lasthandshake:
                multiply.append(86400)

            for i in range(0, len(tm_rev)):
                seconds += int(tm_rev[i]) * multiply[i]

        except:
            seconds = -1
        gauge_wg_handshake.labels(peer=peer['name'], interface=peer['interface']).set(seconds)
        if (seconds == -1) or (seconds > 300):
            gauge_wg_active.labels(peer=peer['name'], interface=peer['interface']).set(0)
        else:
            gauge_wg_active.labels(peer=peer['name'], interface=peer['interface']).set(1)

def get_mangle():

    list = api.get_resource('/ip/firewall/mangle')
    rules = list.get()
    for rule in rules:
        seconds = 0
        if "Download" in rule['comment']:
            gauge_mangle_rx.labels(comment=rule['comment']).set(int(rule['bytes']))
        elif "Upload" in rule['comment']:
            gauge_mangle_tx.labels(comment=rule['comment']).set(int(rule['bytes']))
        else:
            pass


if __name__ == '__main__':
    start_http_server(8000)

    while True:
        get_wg()
        get_mangle()
        time.sleep(5)