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 -> [PortScanResult] {
 24        await withTaskGroup(of: PortScanResult.self, returning: [PortScanResult].self) { group in
 25            for info in ports {
 26                group.addTask {
 27                    let open = await probe(domain: domain, port: info.port)
 28                    return PortScanResult(port: info.port, service: info.service, open: open)
 29                }
 30            }
 31
 32            var results: [PortScanResult] = []
 33            for await result in group {
 34                results.append(result)
 35            }
 36
 37            // Sort by port number
 38            return results.sorted { $0.port < $1.port }
 39        }
 40    }
 41
 42    static func scanPorts(domain: String, ports: [UInt16], timeout: TimeInterval) async -> [PortScanResult] {
 43        await withTaskGroup(of: PortScanResult.self, returning: [PortScanResult].self) { group in
 44            for port in ports {
 45                let service = self.ports.first(where: { $0.port == port })?.service ?? "Custom"
 46                group.addTask {
 47                    let open = await probe(domain: domain, port: port, timeout: timeout)
 48                    return PortScanResult(
 49                        port: port,
 50                        service: service,
 51                        open: open
 52                    )
 53                }
 54            }
 55
 56            var results: [PortScanResult] = []
 57            for await result in group {
 58                results.append(result)
 59            }
 60
 61            return results.sorted { $0.port < $1.port }
 62        }
 63    }
 64
 65    static func grabBanner(host: String, port: UInt16, timeout: TimeInterval = 3.0) async -> String? {
 66        await withCheckedContinuation { continuation in
 67            guard let nwPort = NWEndpoint.Port(rawValue: port) else {
 68                continuation.resume(returning: nil)
 69                return
 70            }
 71
 72            let connection = NWConnection(host: NWEndpoint.Host(host), port: nwPort, using: .tcp)
 73            let context = BannerContext(connection: connection, continuation: continuation)
 74            let queue = DispatchQueue(label: "portscan.banner.\(port)")
 75
 76            connection.stateUpdateHandler = { state in
 77                switch state {
 78                case .ready:
 79                    connection.receive(minimumIncompleteLength: 1, maximumLength: 256) { data, _, _, error in
 80                        guard error == nil,
 81                              let data,
 82                              !data.isEmpty,
 83                              let rawBanner = String(data: data, encoding: .utf8) else {
 84                            context.finish(with: nil)
 85                            return
 86                        }
 87
 88                        let printableBanner = rawBanner.filter { character in
 89                            guard let scalar = character.unicodeScalars.first,
 90                                  character.unicodeScalars.count == 1 else {
 91                                return false
 92                            }
 93                            return (32...126).contains(scalar.value)
 94                        }
 95
 96                        let banner = String(printableBanner.prefix(80))
 97                        context.finish(with: banner.isEmpty ? nil : banner)
 98                    }
 99                case .failed, .cancelled:
100                    context.finish(with: nil)
101                default:
102                    break
103                }
104            }
105
106            connection.start(queue: queue)
107
108            queue.asyncAfter(deadline: .now() + timeout) {
109                context.finish(with: nil)
110            }
111        }
112    }
113
114    private static func probe(domain: String, port: UInt16) async -> Bool {
115        await probe(domain: domain, port: port, timeout: 5)
116    }
117
118    private static func probe(domain: String, port: UInt16, timeout: TimeInterval) async -> Bool {
119        await withCheckedContinuation { continuation in
120            let host = NWEndpoint.Host(domain)
121            let nwPort = NWEndpoint.Port(rawValue: port)!
122            let connection = NWConnection(host: host, port: nwPort, using: .tcp)
123            let context = ProbeContext(connection: connection, continuation: continuation)
124
125            connection.stateUpdateHandler = { state in
126                switch state {
127                case .ready:
128                    context.finish(open: true)
129                case .failed, .cancelled:
130                    context.finish(open: false)
131                default:
132                    break
133                }
134            }
135
136            let queue = DispatchQueue(label: "portscan.\(port)")
137            connection.start(queue: queue)
138
139            queue.asyncAfter(deadline: .now() + timeout) {
140                context.finish(open: false)
141            }
142        }
143    }
144}
145
146private final class ProbeContext: @unchecked Sendable {
147    private let connection: NWConnection
148    private let continuation: CheckedContinuation<Bool, Never>
149    private let lock = NSLock()
150    private nonisolated(unsafe) var resumed = false
151
152    init(connection: NWConnection, continuation: CheckedContinuation<Bool, Never>) {
153        self.connection = connection
154        self.continuation = continuation
155    }
156
157    nonisolated func finish(open: Bool) {
158        lock.lock()
159        guard !resumed else {
160            lock.unlock()
161            return
162        }
163        resumed = true
164        lock.unlock()
165
166        connection.cancel()
167        continuation.resume(returning: open)
168    }
169}
170
171private final class BannerContext: @unchecked Sendable {
172    private let connection: NWConnection
173    private let continuation: CheckedContinuation<String?, Never>
174    private let lock = NSLock()
175    private nonisolated(unsafe) var resumed = false
176
177    init(connection: NWConnection, continuation: CheckedContinuation<String?, Never>) {
178        self.connection = connection
179        self.continuation = continuation
180    }
181
182    nonisolated func finish(with banner: String?) {
183        lock.lock()
184        guard !resumed else {
185            lock.unlock()
186            return
187        }
188        resumed = true
189        lock.unlock()
190
191        connection.cancel()
192        continuation.resume(returning: banner)
193    }
194}