eulaurarien/export.py

#!/usr/bin/env python3

import database
import argparse
import sys


if __name__ == "__main__":

    # Parsing arguments
    parser = argparse.ArgumentParser(
        description="Export the hostnames rules stored " "in the Database as plain text"
    )
    parser.add_argument(
        "-o",
        "--output",
        type=argparse.FileType("w"),
        default=sys.stdout,
        help="Output file, one rule per line",
    )
    parser.add_argument(
        "-f",
        "--first-party",
        action="store_true",
        help="Only output rules issued from first-party sources",
    )
    parser.add_argument(
        "-e",
        "--end-chain",
        action="store_true",
        help="Only output rules that are not referenced by any other",
    )
    parser.add_argument(
        "-r",
        "--rules",
        action="store_true",
        help="Output all kinds of rules, not just hostnames",
    )
    parser.add_argument(
        "-b",
        "--base-rules",
        action="store_true",
        help="Output base rules "
        "(the ones added by ./feed_rules.py) "
        "(implies --rules)",
    )
    parser.add_argument(
        "-d",
        "--no-dupplicates",
        action="store_true",
        help="Do not output rules that already match a zone/network rule "
        "(e.g. dummy.example.com when there's a zone example.com rule)",
    )
    parser.add_argument(
        "-x",
        "--explain",
        action="store_true",
        help="Show the chain of rules leading to one "
        "(and the number of references they have)",
    )
    parser.add_argument(
        "-c",
        "--count",
        action="store_true",
        help="Show the number of rules per type instead of listing them",
    )
    args = parser.parse_args()

    DB = database.Database()

    if args.count:
        assert not args.explain
        print(
            DB.count_records(
                first_party_only=args.first_party,
                end_chain_only=args.end_chain,
                no_dupplicates=args.no_dupplicates,
                rules_only=args.base_rules,
                hostnames_only=not (args.rules or args.base_rules),
            )
        )
    else:
        for domain in DB.list_records(
            first_party_only=args.first_party,
            end_chain_only=args.end_chain,
            no_dupplicates=args.no_dupplicates,
            rules_only=args.base_rules,
            hostnames_only=not (args.rules or args.base_rules),
            explain=args.explain,
        ):
            print(domain, file=args.output)