1 files changed, 209 insertions, 0 deletions
diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
new file mode 100644
index 0000000..2fa1fd0
--- /dev/null
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -0,0 +1,209 @@
+package com.it_jaros.jscanner;
+
+import java.io.BufferedReader;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.net.*;
+import java.nio.file.Files;
+import java.nio.file.NoSuchFileException;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.BitSet;
+import java.util.List;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.LockSupport;
+import java.util.stream.Stream;
+
+public class Scanner {
+
+    private final Counter counter = new Counter();
+    private final ProgressBar progressBar = new ProgressBar();
+    private final Semaphore socketLimit;
+    private final int timeoutInMillis;
+    private final long delayInNanos;
+    private final int maxWorkersPerHost;
+    private final int maxHostsLimit;
+    private final boolean disableHostCheck;
+
+    public Scanner(int socketLimit, int timeoutInMillis, int delayInMillis, int maxWorkersPerHost, int maxHostsLimit, boolean disableHostCheck) {
+        this.timeoutInMillis = timeoutInMillis;
+        this.delayInNanos = TimeUnit.MILLISECONDS.toNanos(Math.max(0, delayInMillis));
+        this.socketLimit = new Semaphore(socketLimit);
+        this.maxWorkersPerHost = maxWorkersPerHost;
+        this.maxHostsLimit = maxHostsLimit;
+        this.disableHostCheck = disableHostCheck;
+    }
+
+    public List<ScanResult> scan(ScanOptions options) {
+        if (options.hostsFile() != null) {
+            return scanHosts(options.hostsFile(), options.ports());
+        }
+        return scanHosts(options.hostsArgv(), options.ports());
+    }
+
+    public List<ScanResult> scanHosts(String sourceFile, String ports) {
+        try (BufferedReader reader = Files.newBufferedReader(Path.of(sourceFile))) {
+            return scanHosts(reader.lines(), ports);
+        } catch (FileNotFoundException | NoSuchFileException e) {
+            System.out.printf("Could not open file %s%n", sourceFile);
+        } catch (IOException e) {
+            System.out.printf("Error while trying to read hosts: %s -> %s %n", e.getClass().getSimpleName(), e.getMessage());
+        }
+
+        return List.of();
+    }
+
+    public List<ScanResult> scanHosts(List<String> hostsArgv, String ports) {
+        try (Stream<String> stream = hostsArgv.stream()){
+            return scanHosts(stream, ports);
+        } catch (IOException e) {
+            System.out.printf("Error while trying to read hosts: %s -> %s %n", e.getClass().getSimpleName(), e.getMessage());
+        }
+
+        return List.of();
+    }
+
+    public List<ScanResult> scanHosts(Stream<String> hosts, String ports) throws IOException {
+        List<ScanResult> results = new ArrayList<>();
+        progressBar.start();
+        try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor()) {
+            List<Future<ScanResult>> futures = new ArrayList<>();
+            Semaphore maxHosts = new Semaphore(maxHostsLimit);
+            hosts.forEach((final String host) -> {
+                maxHosts.acquireUninterruptibly();
+                futures.add(executor.submit(() -> {
+                    try {
+                        return scanHost(host, ports);
+                    } finally {
+                        maxHosts.release();
+                    }
+                }));
+            });
+
+            futures.forEach((Future<ScanResult> f) -> {
+                try {
+                    results.add(f.get());
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                } catch (ExecutionException e) {
+                    System.out.printf("%s -> %s%n", e.getClass().getSimpleName(), e.getMessage());
+                }
+            });
+        }
+
+        progressBar.stop();
+        System.out.println("--------------------");
+        System.out.println("Finished scan!");
+        System.out.println("Stats:");
+        System.out.println("Peak concurrent connects: " + counter.max());
+        System.out.println("--------------------");
+        return results;
+    }
+
+    private ScanResult scanHost(String host, String ports) {
+        PortRange portRange = new PortRange(ports);
+        AtomicLong portSlotFactory = new AtomicLong(System.nanoTime());
+        Semaphore maxWorkers = new Semaphore(maxWorkersPerHost);
+        BitSet openPorts = new BitSet(PortRange.MAX_PORT);
+        BitSet filteredPorts = new BitSet(PortRange.MAX_PORT);
+
+        Progress progress = new Progress(host, portRange.getTotal(), new AtomicInteger(), new AtomicInteger(), new AtomicInteger());
+        progressBar.submit(progress);
+
+        if(!disableHostCheck && !checkHostOnline(host)) {
+            // Visually show that this host is basically done
+            progress.done().set(progress.total());
+            return new ScanResult(host, openPorts, filteredPorts);
+        }
+
+        try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor()) {
+            List<Future<PortResult>> futures = new ArrayList<>();
+            while (portRange.hasNext()) {
+                maxWorkers.acquireUninterruptibly();
+                final int currentPort = portRange.next();
+                futures.add(executor.submit(() -> {
+                    try {
+                        waitForSlot(portSlotFactory);
+                        PortState state = getPortState(host, currentPort);
+                        if (state.equals(PortState.OPEN)) {
+                            progress.open().incrementAndGet();
+                        } else if (state.equals(PortState.FILTERED)) {
+                            progress.filtered().incrementAndGet();
+                        }
+                        return new PortResult(currentPort, state);
+                    } finally {
+                        maxWorkers.release();
+                        progress.done().incrementAndGet();
+                    }
+                }));
+            }
+
+            List<Throwable> errors = new ArrayList<>();
+            for (Future<PortResult> f : futures) {
+                try {
+                    // bitset is not thread-safe, so it is set
+                    // outside the other virtual threads that update progress
+                    PortResult portResult = f.get();
+                    if (portResult.state().equals(PortState.OPEN)) {
+                        openPorts.set(portResult.port());
+                    } else if (portResult.state().equals(PortState.FILTERED)) {
+                        filteredPorts.set(portResult.port());
+                    }
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                } catch (ExecutionException e) {
+                    errors.add(e);
+                }
+            }
+
+            if (!errors.isEmpty()) {
+                System.out.printf("Errors happened during scan of host %s%nErrors:%s -> %s", host, errors.size(), errors);
+            }
+        }
+
+        return new ScanResult(host, openPorts, filteredPorts);
+    }
+
+    private boolean checkHostOnline(String host) {
+        try {
+            return InetAddress.getByName(host).isReachable(timeoutInMillis);
+        } catch (IOException e) {
+            // something went wrong
+        }
+
+        return false;
+    }
+
+    private void waitForSlot(AtomicLong scanSlotFactory) {
+        if (delayInNanos > 0) {
+            long slot = scanSlotFactory.getAndAdd(delayInNanos);
+            long wait = slot - System.nanoTime();
+            if (wait > 0)
+                LockSupport.parkNanos(wait);
+        }
+    }
+
+    private PortState getPortState(String host, int port) {
+        socketLimit.acquireUninterruptibly();
+        counter.inc();
+        try (Socket socket = new Socket()) {
+            socket.connect(new InetSocketAddress(host, port), timeoutInMillis);
+            return PortState.OPEN;
+        } catch (NoRouteToHostException ignored) {
+            // this can be safely ignored because the port is closed if a host is unreachable
+        } catch (SocketTimeoutException ignored) {
+            return PortState.FILTERED;
+        } catch (ConnectException ignored) {
+            return PortState.CLOSED;
+        } catch (IOException ignored) {
+            // Will happen a lot when scanning for open ports, so not needed
+        }  finally {
+            counter.dec();
+            socketLimit.release();
+        }
+
+        return PortState.CLOSED;
+    }
+}