python-threatexchange/threatexchange/cli/dataset_cmd.py

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

from argparse import ArgumentParser
import collections
import csv
import sys
import typing as t
import logging


from threatexchange.signal_type.signal_base import SignalType
from threatexchange import signal_type
from threatexchange.cli.cli_config import CLISettings
from threatexchange.content_type.content_base import ContentType
from threatexchange.fetcher.collab_config import CollaborationConfigBase
from threatexchange.fetcher.fetch_state import FetchedSignalMetadata

from threatexchange.cli import command_base
from threatexchange import common


class DatasetCommand(command_base.Command):
    """
    Introspect fetched data.

    Can print out contents in simple formats
    (ideal for sending to another system), or regenerate
    index files (ideal if distributing indices for some reason)
    """

    @classmethod
    def init_argparse(cls, settings: CLISettings, ap: ArgumentParser) -> None:
        actions = ap.add_mutually_exclusive_group()
        actions.add_argument(
            "--rebuild-indices",
            "-r",
            action="store_true",
            help="rebuild indices from fetched data",
        )
        actions.add_argument(
            "--clear-indices",
            "-X",
            action="store_true",
            help="clear all indices",
        )
        actions.add_argument(
            "--signal-summary",
            "-S",
            action="store_true",
            help="print summary in terms of signals",
        )
        actions.add_argument(
            "--print-records",
            "-P",
            action="store_true",
            help="print records to screen",
        )

        type_selector = ap.add_mutually_exclusive_group()
        type_selector.add_argument(
            "--only-signals",
            "-s",
            nargs="+",
            default=[],
            type=common.argparse_choices_pre_type(
                choices=[s.get_name() for s in settings.get_all_signal_types()],
                type=settings.get_signal_type,
            ),
            help="only process these sigals",
        )
        type_selector.add_argument(
            "--only-content",
            "-C",
            nargs="+",
            default=[],
            type=common.argparse_choices_pre_type(
                choices=[s.get_name() for s in settings.get_all_content_types()],
                type=settings.get_content_type,
            ),
            help="only process signals for these content types",
        )
        ap.add_argument(
            "--only-collabs",
            "-c",
            nargs="+",
            default=[],
            metavar="NAME",
            help="[-S|-P] only count items with this tag",
        )
        ap.add_argument(
            "--only-tags",
            "-t",
            default=[],
            metavar="STR",
            help="[-S|-P] only count items with these tags",
        )
        ap.add_argument(
            "--signals-only",
            "-i",
            action="store_true",
            help="[-P] only print signals",
        )
        ap.add_argument(
            "--limit",
            "-l",
            action="store_true",
            help="[-P] only print this many records",
        )

    def __init__(
        # These all have defaults to make it easier to call
        # only for rebuld
        self,
        # Mode
        clear_indices: bool = False,
        rebuild_indices: bool = False,
        signal_summary: bool = False,
        print_records: bool = False,
        # Signal selectors
        only_collabs: t.Sequence[str] = (),
        only_signals: t.Sequence[t.Type[SignalType]] = (),
        only_content: t.Sequence[t.Type[ContentType]] = (),
        only_tags: t.Sequence[str] = (),
        # Print stuff
        signals_only: bool = False,
        limit: t.Optional[int] = None,
    ) -> None:
        self.clear_indices = clear_indices
        self.rebuild_indices = rebuild_indices
        self.print_records = print_records
        self.signal_summary = signal_summary

        self.only_collabs = set(only_collabs)
        self.only_signals = set(only_signals)
        self.only_content = set(only_content)
        self.only_tags = set(only_tags)
        self.signals_only = signals_only
        self.limit = limit

    def execute(self, settings: CLISettings) -> None:
        # Maybe consider subcommands?
        if self.clear_indices:
            self.execute_clear_indices(settings)
        elif self.rebuild_indices:
            self.execute_generate_indices(settings)
        elif self.print_records:
            self.execute_print_records(settings)
        elif self.signal_summary:
            self.execute_print_signal_summary(settings)
        else:
            self.execute_print_summary(settings)

    def get_signal_types(self, settings: CLISettings) -> t.Set[t.Type[SignalType]]:
        signal_types = self.only_signals or settings.get_all_signal_types()
        if self.only_content:
            signal_types = [
                s
                for s in signal_types
                if any(c in self.only_content for c in s.get_content_types())
            ]
        return set(signal_types)

    def get_collabs(self, settings: CLISettings) -> t.List[CollaborationConfigBase]:
        collabs = [
            c for c in settings.get_all_collabs(default_to_sample=True) if c.enabled
        ]

        if self.only_collabs:
            collabs = [c for c in collabs if c.name in self.only_collabs]
        return collabs

    def get_signals(
        self, settings: CLISettings, signal_types: t.Iterable[t.Type[SignalType]]
    ) -> t.Dict[
        t.Type[SignalType], t.Dict[str, t.List[t.Tuple[str, FetchedSignalMetadata]]]
    ]:
        collabs = self.get_collabs(settings)

        collab_by_api: t.Dict[str, t.List[CollaborationConfigBase]] = {}
        for collab_config in collabs:
            collab_by_api.setdefault(collab_config.api, []).append(collab_config)
        by_type = {}
        for s_type in signal_types:
            by_signal: t.Dict[
                str,
                t.List[t.Tuple[str, FetchedSignalMetadata]],
            ] = {}
            for collabs_for_store in collab_by_api.values():
                store = settings.get_fetch_store_for_collab(collabs_for_store[0])
                by_collab = store.get_for_signal_type(collabs_for_store, s_type)
                for collab, signals in by_collab.items():
                    for signal, record in signals.items():
                        if self.only_tags:
                            for opinion in record.get_as_opinions():
                                if any(t in self.only_tags for t in opinion.tags):
                                    break
                            else:
                                continue
                        by_signal.setdefault(signal, []).append((collab, record))
            by_type[s_type] = by_signal
        return by_type

    def execute_print_summary(self, settings: CLISettings):
        signals = self.get_signals(settings, self.get_signal_types(settings))
        by_type: t.Dict[str, int] = collections.Counter()
        for s_type, type_signals in signals.items():
            by_type[s_type.get_name()] += len(type_signals)
        for s_name, count in sorted(by_type.items(), key=lambda i: -i[1]):
            self.stderr(f"{s_name}: {count}")

    def execute_print_signal_summary(self, settings):
        raise NotImplementedError
        # signal_types = meta.get_signal_types_by_name()
        # by_signal: t.Dict[str, int] = collections.Counter()
        # for indicator in indicators.values():
        #     for name, signal_type in signal_types.items():
        #         if signal_type.indicator_applies(
        #             indicator.indicator_type, list(indicator.rollup.labels)
        #         ):
        #             by_signal[name] += 1
        # for name, count in sorted(by_signal.items(), key=lambda i: -i[1]):
        #     self.stderr(f"{name}: {count}")

    def execute_print_records(self, settings):
        raise NotImplementedError
        # csv_writer = csv.writer(sys.stdout)
        # for indicator in indicators.values():
        #     if self.indicator_only:
        #         print(indicator.indicator)
        #     else:
        #         csv_writer.writerow(indicator.as_csv_row())

    def execute_clear_indices(self, settings: CLISettings) -> None:
        only_signals = None
        if self.only_signals or self.only_content:
            only_signals = self.get_signal_types(settings)
        settings.index_store.clear(only_signals)

    def execute_generate_indices(self, settings: CLISettings) -> None:
        signal_types = self.get_signal_types(settings)

        for s_type in signal_types:
            index_cls = s_type.get_index_cls()
            signal_by_type = self.get_signals(settings, [s_type])
            signals = signal_by_type.get(s_type, {})
            if not signals:
                logging.info("No signals for %s", s_type.__name__)
                settings.index_store.clear([s_type])
                continue

            self.stderr(
                "Building index for",
                s_type.get_name(),
                f"with {len(signals)} signals...",
            )
            index = index_cls.build(signals.items())
            settings.index_store.store_index(s_type, index)
            self.stderr(f"Index for {s_type.get_name()} ready")