krz/domain-dig

  1import Foundation
  2import Network
  3
  4struct PortScanService {
  5    struct PortInfo: Sendable {
  6        let port: UInt16
  7        let service: String
  8    }
  9
 10    static let ports: [PortInfo] = [
 11        PortInfo(port: 21, service: "FTP"),
 12        PortInfo(port: 22, service: "SSH"),
 13        PortInfo(port: 25, service: "SMTP"),
 14        PortInfo(port: 80, service: "HTTP"),
 15        PortInfo(port: 443, service: "HTTPS"),
 16        PortInfo(port: 587, service: "SMTP (TLS)"),
 17        PortInfo(port: 3306, service: "MySQL"),
 18        PortInfo(port: 5432, service: "PostgreSQL"),
 19        PortInfo(port: 8080, service: "HTTP Alt"),
 20        PortInfo(port: 8443, service: "HTTPS Alt"),
 21    ]
 22
 23    static func scanAll(domain: String) async -> ServiceResult<[PortScanResult]> {
 24        await withTaskGroup(of: PortScanResult.self, returning: [PortScanResult].self) { group in
 25            for info in ports {
 26                group.addTask {
 27                    let result = await probe(domain: domain, port: info.port)
 28                    return PortScanResult(
 29                        port: info.port,
 30                        service: info.service,
 31                        open: result.open,
 32                        kind: .standard,
 33                        durationMs: result.durationMs
 34                    )
 35                }
 36            }
 37
 38            var results: [PortScanResult] = []
 39            for await result in group {
 40                results.append(result)
 41            }
 42
 43            // Sort by port number
 44            let sorted = results.sorted { $0.port < $1.port }
 45            return sorted.isEmpty ? [] : sorted
 46        }
 47        .pipe { $0.isEmpty ? .empty("No port scan results") : .success($0) }
 48    }
 49
 50    static func scanPorts(domain: String, ports: [UInt16], timeout: TimeInterval) async -> ServiceResult<[PortScanResult]> {
 51        await withTaskGroup(of: PortScanResult.self, returning: [PortScanResult].self) { group in
 52            for port in ports {
 53                let service = self.ports.first(where: { $0.port == port })?.service ?? "Custom"
 54                group.addTask {
 55                    let result = await probe(domain: domain, port: port, timeout: timeout)
 56                    return PortScanResult(
 57                        port: port,
 58                        service: service,
 59                        open: result.open,
 60                        kind: .custom,
 61                        durationMs: result.durationMs
 62                    )
 63                }
 64            }
 65
 66            var results: [PortScanResult] = []
 67            for await result in group {
 68                results.append(result)
 69            }
 70
 71            let sorted = results.sorted { $0.port < $1.port }
 72            return sorted.isEmpty ? [] : sorted
 73        }
 74        .pipe { $0.isEmpty ? .empty("No custom port scan results") : .success($0) }
 75    }
 76
 77    static func grabBanner(host: String, port: UInt16, timeout: TimeInterval = 3.0) async -> String? {
 78        await withCheckedContinuation { continuation in
 79            guard let nwPort = NWEndpoint.Port(rawValue: port) else {
 80                continuation.resume(returning: nil)
 81                return
 82            }
 83
 84            let connection = NWConnection(host: NWEndpoint.Host(host), port: nwPort, using: .tcp)
 85            let context = BannerContext(connection: connection, continuation: continuation)
 86            let queue = DispatchQueue(label: "portscan.banner.\(port)")
 87
 88            connection.stateUpdateHandler = { state in
 89                switch state {
 90                case .ready:
 91                    connection.receive(minimumIncompleteLength: 1, maximumLength: 256) { data, _, _, error in
 92                        context.finish(with: printableBanner(from: data, error: error))
 93                    }
 94                case .failed, .cancelled:
 95                    context.finish(with: nil)
 96                default:
 97                    break
 98                }
 99            }
100
101            connection.start(queue: queue)
102
103            queue.asyncAfter(deadline: .now() + timeout) {
104                context.finish(with: nil)
105            }
106        }
107    }
108
109    /// Pure data transformation, called from the connection's background queue —
110    /// `nonisolated` opts it out of the project's MainActor default, which would
111    /// otherwise make this call a data-race diagnostic under Swift 6.
112    private nonisolated static func printableBanner(from data: Data?, error: Error?) -> String? {
113        guard error == nil,
114              let data,
115              !data.isEmpty,
116              let rawBanner = String(data: data, encoding: .utf8) else {
117            return nil
118        }
119
120        let printable = rawBanner.filter { character in
121            guard let scalar = character.unicodeScalars.first,
122                  character.unicodeScalars.count == 1 else {
123                return false
124            }
125            return (32...126).contains(scalar.value)
126        }
127
128        let banner = String(printable.prefix(80))
129        return banner.isEmpty ? nil : banner
130    }
131
132    private static func probe(domain: String, port: UInt16) async -> PortProbeResult {
133        await probe(domain: domain, port: port, timeout: 1.5)
134    }
135
136    private static func probe(domain: String, port: UInt16, timeout: TimeInterval) async -> PortProbeResult {
137        await withCheckedContinuation { continuation in
138            let host = NWEndpoint.Host(domain)
139            let nwPort = NWEndpoint.Port(rawValue: port)!
140            let connection = NWConnection(host: host, port: nwPort, using: .tcp)
141            let context = ProbeContext(connection: connection, continuation: continuation)
142
143            connection.stateUpdateHandler = { state in
144                switch state {
145                case .ready:
146                    context.finish(open: true)
147                case .failed, .cancelled:
148                    context.finish(open: false)
149                default:
150                    break
151                }
152            }
153
154            let queue = DispatchQueue(label: "portscan.\(port)")
155            connection.start(queue: queue)
156
157            queue.asyncAfter(deadline: .now() + timeout) {
158                context.finish(open: false)
159            }
160        }
161    }
162}
163
164private struct PortProbeResult: Sendable {
165    let open: Bool
166    let durationMs: Int?
167}
168
169private final class ProbeContext: @unchecked Sendable {
170    private let connection: NWConnection
171    private let continuation: CheckedContinuation<PortProbeResult, Never>
172    private let start = CFAbsoluteTimeGetCurrent()
173    private let lock = NSLock()
174    private nonisolated(unsafe) var resumed = false
175
176    init(connection: NWConnection, continuation: CheckedContinuation<PortProbeResult, Never>) {
177        self.connection = connection
178        self.continuation = continuation
179    }
180
181    nonisolated func finish(open: Bool) {
182        lock.lock()
183        guard !resumed else {
184            lock.unlock()
185            return
186        }
187        resumed = true
188        lock.unlock()
189
190        connection.cancel()
191        let elapsedMs = Int((CFAbsoluteTimeGetCurrent() - start) * 1000)
192        continuation.resume(returning: PortProbeResult(
193            open: open,
194            durationMs: elapsedMs >= 0 ? elapsedMs : nil
195        ))
196    }
197}
198
199private extension Array {
200    func pipe<T>(_ transform: (Self) -> T) -> T {
201        transform(self)
202    }
203}
204
205private final class BannerContext: @unchecked Sendable {
206    private let connection: NWConnection
207    private let continuation: CheckedContinuation<String?, Never>
208    private let lock = NSLock()
209    private nonisolated(unsafe) var resumed = false
210
211    init(connection: NWConnection, continuation: CheckedContinuation<String?, Never>) {
212        self.connection = connection
213        self.continuation = continuation
214    }
215
216    nonisolated func finish(with banner: String?) {
217        lock.lock()
218        guard !resumed else {
219            lock.unlock()
220            return
221        }
222        resumed = true
223        lock.unlock()
224
225        connection.cancel()
226        continuation.resume(returning: banner)
227    }
228}