kindwolf.org Git repositories woobank-utils / master bin / woobank-graphite
master

Tree @master (Download .tar.gz)

woobank-graphite @masterraw · history · blame

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright © 2017-2022 Xavier G. <xavier.woobank@kindwolf.org>
# This work is free. You can redistribute it and/or modify it under the
# terms of the Do What The Fuck You Want To Public License, Version 2,
# as published by Sam Hocevar. See the COPYING file for more details.
"""
woobank-graphite takes a JSON file, as created by woobank-cacher when running
and caching a "woob bank list" command, and sends it to a Graphite instance.
"""

import sys
import json
import time
import socket
import argparse
from woobank_utils import DEFAULT_CURRENCY


def main():
    """
    Main function.
    """
    args = parse_args()
    accounts = json.load(args.json_file)
    timestamp = time.time()
    graphite_socket = graphite_connect(args.host, args.port)
    for account in accounts.get('data', []):
        send_account(graphite_socket, account, timestamp)
    graphite_socket.close()

def parse_args():
    """
    Parse command-line arguments.
    """
    description = 'Send the contents of a woobank-cacher JSON file to Graphite.'
    args = argparse.ArgumentParser(description=description)
    args.add_argument('-gh', '--graphite-host', dest='host', default='::1',
                      help='Graphite host; defaults to ::1.')
    args.add_argument('-gp', '--graphite-port', dest='port', default='2003', type=int,
                      help='Graphite port; defaults to 2003.')
    args.add_argument('json_file', nargs='?', type=argparse.FileType('r'), default=sys.stdin,
                      help='JSON file to send to Graphite; defaults to "-" (standard input).')
    return args.parse_args()

def graphite_connect(host, port):
    """
    Connect to the Graphite instance listening on host and port.
    """
    addr_infos = socket.getaddrinfo(host, port, 0, 0, socket.IPPROTO_TCP)
    _, _, _, _, sockaddr = addr_infos[0]
    graphite_socket = socket.socket(socket.AF_INET6, socket.SOCK_STREAM, 0)
    graphite_socket.connect(sockaddr)
    return graphite_socket

def get_path(account):
    """
    Return the path to send to Graphite for the given bank account.
    The following scheme is used:
    bank.<bank id>.<account id>-balance-<currency>
    """
    (account_id, bank) = account['id'].split('@')
    currency = account.get('currency', DEFAULT_CURRENCY).lower()
    return 'bank.%s.%s-balance-%s' % (bank, account_id, currency)

def send_account(graphite_socket, account, timestamp):
    """
    Send the balance of the given bank account to Graphite over the given TCP
    socket.
    """
    pattern = '%s %.2f %d\n'
    values = (get_path(account), float(account['balance']), int(timestamp))
    line = pattern % values
    graphite_socket.send(line.encode('utf-8'))

if __name__ == '__main__':
    main()