1 files changed, 102 insertions, 1 deletions
diff --git a/DomainDig/PortScanService.swift b/DomainDig/PortScanService.swift
index 63d0571..4dc9001 100644
--- a/DomainDig/PortScanService.swift
+++ b/DomainDig/PortScanService.swift
@@ -39,7 +39,83 @@ struct PortScanService {
         }
     }
 
+    static func scanPorts(domain: String, ports: [UInt16], timeout: TimeInterval) async -> [PortScanResult] {
+        await withTaskGroup(of: PortScanResult.self, returning: [PortScanResult].self) { group in
+            for port in ports {
+                let service = self.ports.first(where: { $0.port == port })?.service ?? "Custom"
+                group.addTask {
+                    let open = await probe(domain: domain, port: port, timeout: timeout)
+                    return PortScanResult(
+                        port: port,
+                        service: service,
+                        open: open
+                    )
+                }
+            }
+
+            var results: [PortScanResult] = []
+            for await result in group {
+                results.append(result)
+            }
+
+            return results.sorted { $0.port < $1.port }
+        }
+    }
+
+    static func grabBanner(host: String, port: UInt16, timeout: TimeInterval = 3.0) async -> String? {
+        await withCheckedContinuation { continuation in
+            guard let nwPort = NWEndpoint.Port(rawValue: port) else {
+                continuation.resume(returning: nil)
+                return
+            }
+
+            let connection = NWConnection(host: NWEndpoint.Host(host), port: nwPort, using: .tcp)
+            let context = BannerContext(connection: connection, continuation: continuation)
+            let queue = DispatchQueue(label: "portscan.banner.\(port)")
+
+            connection.stateUpdateHandler = { state in
+                switch state {
+                case .ready:
+                    connection.receive(minimumIncompleteLength: 1, maximumLength: 256) { data, _, _, error in
+                        guard error == nil,
+                              let data,
+                              !data.isEmpty,
+                              let rawBanner = String(data: data, encoding: .utf8) else {
+                            context.finish(with: nil)
+                            return
+                        }
+
+                        let printableBanner = rawBanner.filter { character in
+                            guard let scalar = character.unicodeScalars.first,
+                                  character.unicodeScalars.count == 1 else {
+                                return false
+                            }
+                            return (32...126).contains(scalar.value)
+                        }
+
+                        let banner = String(printableBanner.prefix(80))
+                        context.finish(with: banner.isEmpty ? nil : banner)
+                    }
+                case .failed, .cancelled:
+                    context.finish(with: nil)
+                default:
+                    break
+                }
+            }
+
+            connection.start(queue: queue)
+
+            queue.asyncAfter(deadline: .now() + timeout) {
+                context.finish(with: nil)
+            }
+        }
+    }
+
     private static func probe(domain: String, port: UInt16) async -> Bool {
+        await probe(domain: domain, port: port, timeout: 5)
+    }
+
+    private static func probe(domain: String, port: UInt16, timeout: TimeInterval) async -> Bool {
         await withCheckedContinuation { continuation in
             let host = NWEndpoint.Host(domain)
             let nwPort = NWEndpoint.Port(rawValue: port)!
@@ -60,7 +136,7 @@ struct PortScanService {
             let queue = DispatchQueue(label: "portscan.\(port)")
             connection.start(queue: queue)
 
-            queue.asyncAfter(deadline: .now() + 3) {
+            queue.asyncAfter(deadline: .now() + timeout) {
                 context.finish(open: false)
             }
         }
@@ -91,3 +167,28 @@ private final class ProbeContext: @unchecked Sendable {
         continuation.resume(returning: open)
     }
 }
+
+private final class BannerContext: @unchecked Sendable {
+    private let connection: NWConnection
+    private let continuation: CheckedContinuation<String?, Never>
+    private let lock = NSLock()
+    private nonisolated(unsafe) var resumed = false
+
+    init(connection: NWConnection, continuation: CheckedContinuation<String?, Never>) {
+        self.connection = connection
+        self.continuation = continuation
+    }
+
+    nonisolated func finish(with banner: String?) {
+        lock.lock()
+        guard !resumed else {
+            lock.unlock()
+            return
+        }
+        resumed = true
+        lock.unlock()
+
+        connection.cancel()
+        continuation.resume(returning: banner)
+    }
+}