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

Tree @master (Download .tar.gz)

woobank-cacher @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-cacher runs a list of "woob bank" commands defined in its configuration
file. For each command, it has the ability to inject simple, additional data to
the result before saving it to a JSON file.
That JSON file can then act as a local cache, which other tools can leverage
for various purposes without emitting too many requests to bank websites.
"""

import os
import re
import json
import copy
import time
import argparse
import subprocess
from woobank_utils import mkdir, DEFAULT_CONF_DIR

def get_conf(conf_path):
    """
    Return the contents of the woobank-cacher configuration.
    """
    conf_path = os.path.expanduser(conf_path)
    with open(conf_path, 'r') as filedesc:
        return json.load(filedesc)

def run_woob_bank(conf):
    """
    Pipe the commands found in the given woobank-cacher configuration to a
    single "woob bank" process and return their results.
    """
    results = copy.deepcopy(conf)
    # Ensure there is at least one command to run:
    commands = results.get('commands', [])
    if not commands:
        return results
    woob_path = os.path.expanduser(conf.get('woob_path', 'woob'))
    woob_bank_command = [woob_path, 'bank', '--formatter=json']
    woob_bank_process = subprocess.Popen(woob_bank_command, stdin=subprocess.PIPE,
                                       stdout=subprocess.PIPE, close_fds=True, bufsize=-1)
    # Send all commands at once:
    for command in commands:
        command_line = command['command'] + "\n"
        woob_bank_process.stdin.write(command_line.encode('utf-8'))
    woob_bank_process.stdin.write("quit\n".encode('utf-8'))
    woob_bank_process.stdin.flush()
    woob_bank_process.stdin.close()
    # Parse output to extract relevant lines only:
    out_time = time.time()
    out_lines = []
    while True:
        out_line = woob_bank_process.stdout.readline().decode('utf-8')
        if not out_line:
            break
        rem = re.search(r'^bank>\s*(.+)$', out_line)
        if rem:
            out_lines.append(rem.group(1))
    for i in range(len(commands)):
        out_data = json.loads(out_lines[i])
        commands[i]['result'] = {'timestamp': out_time, 'data': out_data}
    return results

def process_additions(results):
    """
    Inject additional data described in the woobank-cacher configuration to the
    given results.
    """
    for command in results.get('commands', []):
        additions = command.get('additions', {})
        result = command.get('result', {})
        if not additions or not result or 'data' not in result:
            continue

        processed_items = []
        for item in result.get('data', []):
            processed_items.append(item)
            item_id = item.get('id')
            if item_id is not None and item_id in additions:
                processed_items.append(additions.get(item_id))
        result['data'] = processed_items
    return results

def write_cache(results):
    """
    Write the given results to cache files as described in the woobank-cacher
    configuration.
    """
    for command in results.get('commands', []):
        if 'result' not in command or 'cache_path' not in command:
            continue
		# There is no point caching data-less results:
        if 'data' not in command['result'] or not command['result']['data']:
            continue
        cache_path = os.path.expanduser(command['cache_path'])
        mkdir(os.path.dirname(cache_path), 0o700)
        with open(cache_path, 'w') as filedesc:
            json.dump(command['result'], filedesc, indent=4)

def parse_args():
    """
    Parse command-line arguments.
    """
    description = 'Cache the output of woob bank commands to JSON files.'
    args_parser = argparse.ArgumentParser(description=description)
    default_conf = os.path.join(DEFAULT_CONF_DIR, 'cacher.json')
    args_parser.add_argument('-c', '--config', default=default_conf,
                             help='configuration file to use')
    return args_parser.parse_args()

def main():
    """
    Main function.
    """
    args = parse_args()
    conf = get_conf(args.config)
    results = run_woob_bank(conf)
    process_additions(results)
    write_cache(results)

if __name__ == '__main__':
    main()