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                        guard error == nil,
 93                              let data,
 94                              !data.isEmpty,
 95                              let rawBanner = String(data: data, encoding: .utf8) else {
 96                            context.finish(with: nil)
 97                            return
 98                        }
 99
100                        let printableBanner = rawBanner.filter { character in
101                            guard let scalar = character.unicodeScalars.first,
102                                  character.unicodeScalars.count == 1 else {
103                                return false
104                            }
105                            return (32...126).contains(scalar.value)
106                        }
107
108                        let banner = String(printableBanner.prefix(80))
109                        context.finish(with: banner.isEmpty ? nil : banner)
110                    }
111                case .failed, .cancelled:
112                    context.finish(with: nil)
113                default:
114                    break
115                }
116            }
117
118            connection.start(queue: queue)
119
120            queue.asyncAfter(deadline: .now() + timeout) {
121                context.finish(with: nil)
122            }
123        }
124    }
125
126    private static func probe(domain: String, port: UInt16) async -> PortProbeResult {
127        await probe(domain: domain, port: port, timeout: 1.5)
128    }
129
130    private static func probe(domain: String, port: UInt16, timeout: TimeInterval) async -> PortProbeResult {
131        await withCheckedContinuation { continuation in
132            let host = NWEndpoint.Host(domain)
133            let nwPort = NWEndpoint.Port(rawValue: port)!
134            let connection = NWConnection(host: host, port: nwPort, using: .tcp)
135            let context = ProbeContext(connection: connection, continuation: continuation)
136
137            connection.stateUpdateHandler = { state in
138                switch state {
139                case .ready:
140                    context.finish(open: true)
141                case .failed, .cancelled:
142                    context.finish(open: false)
143                default:
144                    break
145                }
146            }
147
148            let queue = DispatchQueue(label: "portscan.\(port)")
149            connection.start(queue: queue)
150
151            queue.asyncAfter(deadline: .now() + timeout) {
152                context.finish(open: false)
153            }
154        }
155    }
156}
157
158private struct PortProbeResult: Sendable {
159    let open: Bool
160    let durationMs: Int?
161}
162
163private final class ProbeContext: @unchecked Sendable {
164    private let connection: NWConnection
165    private let continuation: CheckedContinuation<PortProbeResult, Never>
166    private let start = CFAbsoluteTimeGetCurrent()
167    private let lock = NSLock()
168    private nonisolated(unsafe) var resumed = false
169
170    init(connection: NWConnection, continuation: CheckedContinuation<PortProbeResult, Never>) {
171        self.connection = connection
172        self.continuation = continuation
173    }
174
175    nonisolated func finish(open: Bool) {
176        lock.lock()
177        guard !resumed else {
178            lock.unlock()
179            return
180        }
181        resumed = true
182        lock.unlock()
183
184        connection.cancel()
185        let elapsedMs = Int((CFAbsoluteTimeGetCurrent() - start) * 1000)
186        continuation.resume(returning: PortProbeResult(
187            open: open,
188            durationMs: elapsedMs >= 0 ? elapsedMs : nil
189        ))
190    }
191}
192
193private extension Array {
194    func pipe<T>(_ transform: (Self) -> T) -> T {
195        transform(self)
196    }
197}
198
199private final class BannerContext: @unchecked Sendable {
200    private let connection: NWConnection
201    private let continuation: CheckedContinuation<String?, Never>
202    private let lock = NSLock()
203    private nonisolated(unsafe) var resumed = false
204
205    init(connection: NWConnection, continuation: CheckedContinuation<String?, Never>) {
206        self.connection = connection
207        self.continuation = continuation
208    }
209
210    nonisolated func finish(with banner: String?) {
211        lock.lock()
212        guard !resumed else {
213            lock.unlock()
214            return
215        }
216        resumed = true
217        lock.unlock()
218
219        connection.cancel()
220        continuation.resume(returning: banner)
221    }
222}