1 files changed, 37 insertions, 18 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index 2fa1fd0..ca5aa4e 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -3,12 +3,15 @@ package com.it_jaros.jscanner;
 import java.io.BufferedReader;
 import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.io.InputStream;
 import java.net.*;
+import java.nio.charset.StandardCharsets;
 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.HashMap;
 import java.util.List;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -108,6 +111,7 @@ public class Scanner {
         Semaphore maxWorkers = new Semaphore(maxWorkersPerHost);
         BitSet openPorts = new BitSet(PortRange.MAX_PORT);
         BitSet filteredPorts = new BitSet(PortRange.MAX_PORT);
+        HashMap<Integer, ServiceType> bannerRecognition = new HashMap<Integer, ServiceType>();
 
         Progress progress = new Progress(host, portRange.getTotal(), new AtomicInteger(), new AtomicInteger(), new AtomicInteger());
         progressBar.submit(progress);
@@ -115,7 +119,7 @@ public class Scanner {
         if(!disableHostCheck && !checkHostOnline(host)) {
             // Visually show that this host is basically done
             progress.done().set(progress.total());
-            return new ScanResult(host, openPorts, filteredPorts);
+            return new ScanResult(host, openPorts, filteredPorts, bannerRecognition);
         }
 
         try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor()) {
@@ -126,13 +130,7 @@ public class Scanner {
                 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);
+                        return checkPort(host, currentPort);
                     } finally {
                         maxWorkers.release();
                         progress.done().incrementAndGet();
@@ -146,10 +144,11 @@ public class Scanner {
                     // 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());
+                    if (portResult.getState().equals(PortState.OPEN)) {
+                        openPorts.set(portResult.getPort());
+                        bannerRecognition.put(portResult.getPort(), ServiceDetector.detect(portResult.getBanner()));
+                    } else if (portResult.getState().equals(PortState.FILTERED)) {
+                        filteredPorts.set(portResult.getPort());
                     }
                 } catch (InterruptedException e) {
                     Thread.currentThread().interrupt();
@@ -163,7 +162,7 @@ public class Scanner {
             }
         }
 
-        return new ScanResult(host, openPorts, filteredPorts);
+        return new ScanResult(host, openPorts, filteredPorts, bannerRecognition);
     }
 
     private boolean checkHostOnline(String host) {
@@ -185,18 +184,22 @@ public class Scanner {
         }
     }
 
-    private PortState getPortState(String host, int port) {
+    private PortResult checkPort(String host, int port) {
         socketLimit.acquireUninterruptibly();
         counter.inc();
+        PortResult result = new PortResult();
+        result.setPort(port);
+        result.setState(PortState.UNKNOWN);
         try (Socket socket = new Socket()) {
             socket.connect(new InetSocketAddress(host, port), timeoutInMillis);
-            return PortState.OPEN;
+            result.setBanner(getBanner(socket));
+            result.setState(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;
+            result.setState(PortState.FILTERED);
         } catch (ConnectException ignored) {
-            return PortState.CLOSED;
+            result.setState(PortState.CLOSED);
         } catch (IOException ignored) {
             // Will happen a lot when scanning for open ports, so not needed
         }  finally {
@@ -204,6 +207,22 @@ public class Scanner {
             socketLimit.release();
         }
 
-        return PortState.CLOSED;
+        return result;
+    }
+
+    private String getBanner(Socket socket) {
+        int READ_BUFFER_SIZE = 1024;
+        byte[] buffer = new byte[READ_BUFFER_SIZE];
+        try(InputStream input = socket.getInputStream()) {
+            socket.setSoTimeout(timeoutInMillis);
+            int bytesRead = input.read(buffer);
+            if (bytesRead <= 0) {
+                return null;
+            }
+            return new String(buffer,0, bytesRead, StandardCharsets.UTF_8).trim();
+        } catch (IOException e) {
+            // we ignore this failure
+        }
+        return null;
     }
 }