package com.pluralsight.protectivetechnology.kafka;

import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.common.config.SslConfigs;
import org.apache.kafka.common.serialization.Serdes;
import org.apache.kafka.streams.*;
import org.apache.kafka.streams.kstream.KStream;
import org.apache.kafka.streams.kstream.TimeWindows;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.Arrays;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class NetworkLogsAnalyzer {

    private static final Logger log = LoggerFactory.getLogger(NetworkLogsAnalyzer.class);

    public static final Pattern SOURCE_PATTERN = Pattern.compile("IP (.*?) >");
    public static final Pattern PORT_PATTERN = Pattern.compile("> [\\d]+\\.[\\d]+\\.[\\d]+\\.[\\d]+\\.(.*?):");

    public static final String NETWORK_LOGS_TOPIC = "network-logs";
    public static final String PORT_SCAN_ATTACKS_ALERTS_TOPIC = "alerts.port-scan-attack";

    public static void main(String[] args) throws ExecutionException, InterruptedException {
        Properties props = new Properties();
        props.put(StreamsConfig.APPLICATION_ID_CONFIG, "network.logs.analyzer.app");
        props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "broker-1:9091,broker-2:9092,broker-3:9093");
        props.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.String().getClass());
        props.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.String().getClass());

        props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, "SSL");
        props.put(SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, "/etc/kafka/secrets/kstreams.keystore.jks");
        props.put(SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, "notsecret");
        props.put(SslConfigs.SSL_KEY_PASSWORD_CONFIG, "notsecret");
        props.put(SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, "/etc/kafka/secrets/kstreams.truststore.jks");
        props.put(SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, "notsecret");

        StreamsBuilder builder = new StreamsBuilder();

        KStream<Long, String> stream = builder.stream(NETWORK_LOGS_TOPIC);
        stream
                .peek((key, value) -> log.info("Consumed record: {}, {}", key, value))
                .map((key, value) -> KeyValue.pair(extractSource(value), extractPort(value)))
                .filter((key, value) -> !key.isBlank() && !value.isBlank())
                .peek((key, value) -> log.info("Mapped record: {}, {}", key, value))
                .groupByKey()
                .windowedBy(TimeWindows.ofSizeWithNoGrace(Duration.ofSeconds(10)))
                .reduce((port1, port2) -> port1 + "," + port2)
                .toStream()
                .peek((key, value) -> log.info("Streamed record: {}, {}", key, value))
                .mapValues(ports -> Set.copyOf(Arrays.asList(ports.split(",")))) // set of devices
                .peek((key, value) -> log.info("Intermediate record: {}, {}", key.key(), value))
                .mapValues(Set::size)
                .peek((key, value) -> log.info("Final record: {}, {}", key.key(), value))
                .filter((source, count) -> count > 10)
                .map((key, value) -> KeyValue.pair(key.key(), value.toString()))
                .peek((key, value) -> log.info("Port Scan Attack Detected from source {}", key))
                .to(PORT_SCAN_ATTACKS_ALERTS_TOPIC);

        Topology topology = builder.build();
        KafkaStreams streams = new KafkaStreams(topology, props);

        Thread haltedHook = new Thread(streams::close);
        Runtime.getRuntime().addShutdownHook(haltedHook);
        streams.start();
    }

    private static String extractSource(String log) {
        Matcher matcher = SOURCE_PATTERN.matcher(log);
        boolean found = matcher.find();
        return found ? matcher.group(1) : "";
    }

    private static String extractPort(String log) {
        Matcher matcher = PORT_PATTERN.matcher(log);
        boolean found = matcher.find();
        return found ? matcher.group(1) : "";
    }
}
