path: root/src/main/java/com/it_jaros/jscanner/ProgressBar.java

package com.it_jaros.jscanner;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.BitSet;
import java.util.Map;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

public class ProgressBar {

    private final ScheduledExecutorService ui = Executors.newSingleThreadScheduledExecutor();
    private final Scan scan;
    private final boolean quiet;
    private final Object outputLock = new Object();

    public ProgressBar(Scan scan, boolean quiet) {
        this.scan = scan;
        this.quiet = quiet;
    }

    private void clearStatusLine() {
        System.err.print("\r\033[2K");
        System.err.flush();
    }

    public void printResult(ScanResult result, boolean showFilteredPorts) {
        synchronized (outputLock) {
            if (result.openPorts().isEmpty()
                    && (!showFilteredPorts || result.filteredPorts().isEmpty())
                    && result.errors().isEmpty()) {
                return;
            }

            clearStatusLine();
            printErrors(result);
            printPortsResults(result, showFilteredPorts);
            printStatusLine();
        }
    }

    private void printErrors(ScanResult result) {
        if (result.errors().isEmpty()) return;

        System.err.printf(
                "error: host scan failed | host=%s | errors: %d%n",
                result.host(),
                result.errors().size()
        );

        for (ScanFailure error : result.errors()) {
            System.err.printf(
                    "port: %s: %s %s%n",
                    error.port(),
                    error.exception().type(),
                    error.exception().message()
            );
        }
    }

    private void printPortsResults(ScanResult result, boolean showFilteredPorts) {
        if (result.openPorts().isEmpty() && (!showFilteredPorts || result.filteredPorts().isEmpty())) {
            return;
        }

        System.out.println(result.host());
        if (!result.openPorts().isEmpty()) {
            System.out.printf("\t (%d) open:\t%s%n", result.openPorts().cardinality(), map(result.openPorts(), result.bannerRecognition()));
        }
        if (showFilteredPorts && !result.filteredPorts().isEmpty()) {
            System.out.printf("\t (%d) filtered:\t%s%n", result.filteredPorts().cardinality(), map(result.filteredPorts()));
        }
        System.out.println();
        System.out.flush();
    }

    private static String map(BitSet ports) {
        return ports.stream().mapToObj(String::valueOf).collect(Collectors.joining(","));
    }

    private static String map(BitSet ports, Map<Integer, ServiceType> bannerRecognition) {
        return ports.stream().mapToObj(port -> {
            ServiceType serviceType = bannerRecognition.getOrDefault(port, ServiceType.UNKNOWN);

            if (serviceType == ServiceType.UNKNOWN) {
                return String.valueOf(port);
            }

            return port + "/" + serviceType.name().toLowerCase();
        }).collect(Collectors.joining(","));
    }

    private static int getColumnWidth() {
        try {
            Process process = new ProcessBuilder("sh", "-c", "stty size < /dev/tty").start();

            try (BufferedReader reader = new BufferedReader(
                    new InputStreamReader(process.getInputStream()))) {
                String line = reader.readLine();
                if (line != null && !line.isBlank()) {
                    Matcher matcher = Pattern.compile("^\\d+\\s+(\\d+)").matcher(line.trim());
                    if (matcher.matches()) {
                        return Integer.parseInt(matcher.group(1));
                    }
                }
            }
        } catch (Exception ignored) {
            // does not matter why it did not work
        }

        return 120;
    }

    public void start() {
        if (quiet) {
            return;
        }

        // although we don't print anything, we still need to
        // gc all the done hosts
        System.err.print("Scanning targets...\n\n");
        ui.scheduleAtFixedRate(() -> {
            synchronized (outputLock) {
                printStatusLine();
            }
        }, 1, 1000, TimeUnit.MILLISECONDS);
    }

    public void stop() {
        ui.shutdownNow();
        try {
            ui.awaitTermination(1, TimeUnit.SECONDS);
        } catch (InterruptedException ignored) {
            Thread.currentThread().interrupt();
        }
        clearStatusLine();
        printStats();
    }

    private void printStats() {
        Scan.Statistics stats = scan.getStatistics();
        System.err.println("\n--------------------");
        System.err.println("Stats:");
        System.err.println("Total hosts: " + stats.hostTotal());
        System.err.println("Total ports: " + stats.portTotal());
        System.err.println("Duration: " + formatDuration(stats.durationInMillis()));
        System.err.println("Peak concurrent hosts: " + stats.hostMaxConcurrent());
        System.err.println("Peak concurrent sockets: " + stats.socketMaxConcurrent());
        System.err.println("Peak concurrent workers: " + stats.threadMaxConcurrent());
        System.err.println("--------------------\n");
    }

    private static String formatDuration(long millis) {
        if (millis <= 0) return "";
        long seconds = millis / 1000 % 60;
        long minutes = millis / 60000 % 60;
        long hours = millis / 3600000 % 24;
        long days = millis / 86400000;

        if (days > 0) return String.format("%dd %dh:%dm:%ds", days, hours, minutes, seconds);
        if (hours > 0) return String.format("%dh:%dm:%ds", hours, minutes, seconds);
        if (minutes > 0) return String.format("%dm:%ds", minutes, seconds);
        return String.format("%ds", seconds);
    }


    private void printStatusLine() {
        Scan.Statistics stats = scan.getStatistics();
        int total = stats.hostTotal();
        int running = stats.hostCurrent();
        int done = total - running;
        long percent = total == 0 ? 100 : done * 100L / total;

        String hostStat = String.format(
                "hosts: %d running, %d done",
                running,
                done
        );
        String portStat = String.format(
                " | ports: %d running, %d done",
                stats.portCurrent(),
                stats.portTotal()
        );
        String durationStat = String.format(
                " | %s",
                formatDuration(stats.durationInMillis())
        );
        String socketStat = String.format(
                " | sockets: %d active, %d peak",
                stats.socketCurrent(),
                stats.socketMaxConcurrent()
        );
        String workerStat = String.format(
                " | workers: %d active, %d peak",
                stats.threadCurrent(),
                stats.threadMaxConcurrent()
        );

        int availableWidth = ProgressBar.getColumnWidth();
        String statusLine = appendIfEnoughSpace("", availableWidth, hostStat, portStat, durationStat, socketStat, workerStat);

        String barFormat = " [%s] %3d%%";
        int remainingWidth = availableWidth - statusLine.length() - barFormat.length();
        if (remainingWidth > 0) {
            String progressBar = String.format(
                    barFormat,
                    progressBar(percent, remainingWidth),
                    percent
            );
            statusLine = appendIfEnoughSpace(statusLine, availableWidth, progressBar);
        }

        remainingWidth = availableWidth - statusLine.length();
        if (remainingWidth > 0) {
            statusLine += " ".repeat(remainingWidth);
        }
        System.err.print("\r" + statusLine);
        System.err.flush();
    }

    private String appendIfEnoughSpace(String current, int availableSize, String... additional) {
        StringBuilder builder = new StringBuilder(current);
        for (String add : additional) {
            if ((builder.length() + add.length()) <= availableSize) {
                builder.append(add);
            }
        }

        return builder.toString();
    }

    private static String progressBar(long percent, int width) {
        if (width <= 0) {
            return "";
        }

        int filled = (int) (percent * width / 100);
        return "#".repeat(filled) + "-".repeat(width - filled);
    }
}