fediverse-blocklist-tool/fediverse_blocklist_tool/cli.py

# import tomllib
import argparse
import json
import logging
import requests
import os
import toml

from fediverse_blocklist_tool.models import Instance
from fediverse_blocklist_tool.helpers import blocklist_to_markdown, blocklist_to_toml, blocklist_to_csv, \
    blocklist_to_json


def load_blocklist_file(filename: str) -> [Instance]:
    if filename.endswith("json"):
        with open(filename) as f:
            instances = blocklist_json_to_instances(json.load(f))
    else:
        with open(filename, "r") as f:
            data = toml.load(f)
        instances = []
        for instance_dict in data["instances"]:
            instance = Instance(instance_dict)
            instances.append(instance)
    return instances


def blocklist_json_to_instances(blocklist_json: str) -> [Instance]:
    instances = []
    for i in blocklist_json:
        instances.append(Instance(i))
    return instances


def load_blocklist_from_instance(server: str, token: str) -> [Instance]:
    headers = {
        f'Authorization': f'Bearer {token}',
    }

    response = requests.get(f'https://{server}/api/v1/admin/domain_blocks', headers=headers)
    if response.status_code == 200:
        blocklist_json = json.loads(response.content)
        return blocklist_json_to_instances(blocklist_json)
    else:
        raise ConnectionError(f"Could not connect to the server ({response.status_code}: {response.reason})")


def remove_key_from_dict(dict, key):
    del dict[key]
    return dict


def exporter(blocklist, output=None, format: str = "toml", private: bool = False):
    if format == "toml":
        exported_text = blocklist_to_toml(blocklist, private)
    if format == "csv":
        exported_text = blocklist_to_csv(blocklist, private)
    if format == "markdown":
        exported_text = blocklist_to_markdown(blocklist, private)
    if format == "json":
        exported_text = blocklist_to_json(blocklist, private)

    # Output the text
    if output is not None:
        with open(output, "w") as f:
            f.write(exported_text)
    else:
        print(exported_text)


def merge(input_file, merge_target, format: str = "toml", private: bool = False, overwrite=False):
    """Shows a table in the CLI comparing the local and remote blocklist"""
    from rich.table import Table
    from rich.console import Console
    input_blocklist = load_blocklist_file(input_file)
    merge_target_blocklist = load_blocklist_file(merge_target)
    for input_instance in input_blocklist:
        # If the block is already there with the same parameters we do nothing
        if input_instance in merge_target_blocklist:
            continue
        # Check if there is a domain in the merge target where the input domain is similar
        try:
            merge_target_instance = [merge_target_instance for merge_target_instance in merge_target_blocklist if input_instance.domain == merge_target_instance.domain][0]
            if not overwrite:
                key_input = ""
                while key_input not in ("i", "O"):
                    print(f"Different settings for {input_instance.domain} detected.")
                    Instance.show_diff(input_instance, merge_target_instance)
                    key_input = input("Keep input (i) or original (o) [i/O]")
            elif key_input == "i":
                merge_target_blocklist.append(merge_target_instance)
            else:
                merge_target_blocklist.append(input_instance)
        except KeyError:
            pass


def cli():
    parser = argparse.ArgumentParser(description='Compare, merge, export and deploy blocklist of a fediverse server')
    parser.add_argument('action', choices=['diff', 'deploy', 'export', 'merge'],
                        help="Either use 'diff' to check the difference between local blockĺist and the blocklist on "
                             "the server, 'deploy' to apply the current local blocklist or 'export' to export the remote "
                             "blocklist into a local file. merge can be used to merge a blocklist (given by -i) into "
                             "another (-o)")
    parser.add_argument('-s', '--server', help="The address of the server where you want to deploy (e.g. "
                                               "mastodon.social)")
    parser.add_argument('-t', '--token', help="Authorization token")
    parser.add_argument('-i', '--input-file', help="The blocklist to use")
    parser.add_argument('-r', '--remote-blocklist', help="The remote blocklist as json for debugging reasons")
    parser.add_argument('-o', '--output', help="Filename where to export the blocklist")
    parser.add_argument('-v', '--verbose', action='store_true')
    parser.add_argument('-n', '--no-delete', action='store_true', help="Do not delete existing blocks")
    parser.add_argument('--format', default="toml", type=str, help="Export format: toml|markdown|csv|json")
    parser.add_argument('--private', action='store_true', help="When the flag is set, private comment will also be "
                                                               "exported.")
    args = parser.parse_args()
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.WARN)

    if args.token:
        token = args.token
    else:
        token = os.getenv('MBD_TOKEN')

    """Get a remote blocklist only when necessary"""
    if args.action in ["diff", "deploy", "export"]:
        """if there is a remote blocklist provided load this instead of fetching it from a server (for debugging reasons)"""
        if args.remote_blocklist:
            remote_blocklist = load_blocklist_file(args.remote_blocklist)
        else:
            remote_blocklist = load_blocklist_from_instance(server=args.server, token=token)

    """Load local blocklist only when needed"""
    if args.action in ["diff", "deploy", "merge"]:
        if args.input_file:
            blocklist_filename = args.input_file
        else:
            blocklist_filename = "../blocklist.toml"
        try:
            local_blocklist = load_blocklist_file(blocklist_filename)
        except FileNotFoundError:
            print("Local blocklist file was not found. Make sure to specify it's location via -i")
            exit()

    if args.action == "diff":
        Instance.show_diffs(local_blocklist, remote_blocklist)
    elif args.action == "deploy":
        diffs = Instance.list_diffs(local_blocklist, remote_blocklist)
        Instance.apply_blocks_from_diff(diffs, args.server, token, args.no_delete)
    elif args.action == "export":
        exporter(remote_blocklist, args.output, args.format, args.private)
    elif args.action == "merge":
        merge(args.input_file, args.output, args.format, args.private)


if __name__ == "__main__":
    cli()