armenak / DataDefender

@Getter

@Log4j2

public class ModelDiscoveryConfig {

    @Option(names = { "-l", "--limit" }, description = "Limit discovery to a set number of rows in a table", defaultValue = "1000")

    private Integer limit;

    private List<String> models;

    @Option(names = { "-e", "--extension" }, description = "Adds a call to an extension method "

        + "(e.g. com.strider.datadefender.specialcase.SinDetector.detectSin)")

    private List<String> extensions;

    @Option(names = { "--model-file" }, description = "Adds a custom made opennlp TokenizerME file for data discovery.")

    private List<File> fileModels;

    @Option(names = { "--token-model" }, description = "Override the default built-in token model (English tokens, "

        + "en-token.bin) with a custom token file for use by opennlp's TokenizerModel")

    private File tokenModel;

    @Option(names = { "--probability-threshold" }, description = "Minimum NLP match score to return results for", defaultValue = "0.55")

    private Double probabilityThreshold;

    @Option(names = { "--no-score-calculation" }, description = "If set, includes a column score", negatable = true)

    private Boolean calculateScore = true;

    @Option(names = { "--threshold-count" }, description = "Reports if number of rows found are greater than the defined threshold", defaultValue = "6")

    private Integer thresholdCount;

    @Option(names = { "--threshold-high" }, description = "Reports if number of high risk columns found are greater than the defined threshold", defaultValue = "3")

    private Integer thresholdHighRisk;

    @Option(names = { "-m", "--model" }, description = "Adds a built-in configured opennlp TokenizerME model for data discovery. "

                + "Available models are: ${AVAILABLE-MODELS}")

    public void setModels(List<String> models) {

        Optional<String> unmatched = models.stream().filter((m) -> !Discoverer.BUILT_IN_MODELS.containsKey(m)).findFirst();

        if (unmatched.isPresent()) {

            log.error(

                "A built-in model with the name \"{}\" does not exist. Please specify one of: {}",

                unmatched.get(),

                System.getProperty("AVAILABLE-MODELS")

            );

            throw new IllegalArgumentException("Unmatched built-in model.");

        }

        this.models = models;