##################################################################################
#                                                                                #
# Takes the output from fritz-api.sh , processes it and exports it to Promethues #
# via http on default port 9000 (can be configured via config.yml                #
#                                                                                #
# (C)2022 M.T. Konstapel https://meezenest.nl/mees                               #
#                                                                                # 
# This file is part of fritzbox_exporter.                                        #
#                                                                                #
#    fritzbox_exporter is free software: you can redistribute it and/or modify   #
#    it under the terms of the GNU General Public License as published by        #
#    the Free Software Foundation, either version 3 of the License, or           #
#    (at your option) any later version.                                         #
#                                                                                #
#    fritzbox_exporter is distributed in the hope that it will be useful,        #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of              #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               #
#    GNU General Public License for more details.                                #
#                                                                                #  
#    You should have received a copy of the GNU General Public License           #
#    along with fritzbox_exporter. If not, see <https://www.gnu.org/licenses/>.  #
#                                                                                #
##################################################################################

import time
import random
import subprocess
import json
from os import path
import yaml
from prometheus_client.core import GaugeMetricFamily, REGISTRY, CounterMetricFamily
from prometheus_client import start_http_server
totalRandomNumber = 0

# Collect Fritzbox data for the first time. This way, the json_object is defined globally. This is not the proper way to do it, but it works.
result = subprocess.run(['./fritz-api.sh', '-j'], stdout=subprocess.PIPE)
#convert string to object
json_object = json.loads(result.stdout)

class Fritzbox_collector(object):
    def __init__(self):
        pass
    def collect(self):
        # Connection state
        if json_object["Connection"] in ['Connected']:
            Fritzbox_connection_state = 2
        else:
            Fritzbox_connection_state = 1
        
        gauge = GaugeMetricFamily("fritz_connect_state", "Fritzbox connection state (1=not connected, 2=connected)", labels=["Fritzbox"])
        gauge.add_metric(['1'], Fritzbox_connection_state)
        yield gauge
        
        gauge = GaugeMetricFamily("fritz_down_speed", "Fritzbox downlink speed", labels=["Fritzbox"])
        gauge.add_metric(['1'], json_object["DownstreamSync"])
        yield gauge
        
        gauge = GaugeMetricFamily("fritz_up_speed", "Fritzbox uplink speed", labels=["Fritzbox"])
        gauge.add_metric(['1'], json_object["UpstreamSync"])
        yield gauge
        
        count = CounterMetricFamily("fritz_uptime", "Fritzbox uptime", labels=['Fritzbox'])
        count.add_metric(['1'], json_object["Uptime"])
        yield count
        
        count = CounterMetricFamily("fritz_upload_bw", "Fritzbox upload BW", labels=['Fritzbox'])
        count.add_metric(['1'], json_object["UploadBW"])
        yield count
        
        count = CounterMetricFamily("fritz_download_bw", "Fritzbox download BW", labels=['Fritzbox'])
        count.add_metric(['1'], json_object["DownloadBW"])
        yield count
        
        count = CounterMetricFamily("fritz_total_uploads", "Fritzbox total uploads", labels=['Fritzbox'])
        count.add_metric(['1'], json_object["TotalUploads"])
        yield count
        
        count = CounterMetricFamily("fritz_total_downloads", "Fritzbox total downloads", labels=['Fritzbox'])
        count.add_metric(['1'], json_object["TotalDownloads"])
        yield count        

if __name__ == "__main__":
    port = 9000
    frequency = 15
    if path.exists('config.yml'):
        with open('config.yml', 'r') as config_file:
            try:
                config = yaml.safe_load(config_file)
                port = int(config['port'])
                frequency = config['scrape_frequency']
            except yaml.YAMLError as error:
                print(error)

    start_http_server(port)
    REGISTRY.register(Fritzbox_collector())

    while True: 
        # Collect Fritzbox data
        result = subprocess.run(['./fritz-api.sh', '-j'], stdout=subprocess.PIPE)
        #convert string to json object
        json_object = json.loads(result.stdout)
        
        # period between collection
        time.sleep(frequency)