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    private static func probe(domain: String, port: UInt16) async -> Bool {
43        await withCheckedContinuation { continuation in
44            let host = NWEndpoint.Host(domain)
45            let nwPort = NWEndpoint.Port(rawValue: port)!
46            let connection = NWConnection(host: host, port: nwPort, using: .tcp)
47            let context = ProbeContext(connection: connection, continuation: continuation)
48
49            connection.stateUpdateHandler = { state in
50                switch state {
51                case .ready:
52                    context.finish(open: true)
53                case .failed, .cancelled:
54                    context.finish(open: false)
55                default:
56                    break
57                }
58            }
59
60            let queue = DispatchQueue(label: "portscan.\(port)")
61            connection.start(queue: queue)
62
63            queue.asyncAfter(deadline: .now() + 3) {
64                context.finish(open: false)
65            }
66        }
67    }
68}
69
70private final class ProbeContext: @unchecked Sendable {
71    private let connection: NWConnection
72    private let continuation: CheckedContinuation<Bool, Never>
73    private let lock = NSLock()
74    private nonisolated(unsafe) var resumed = false
75
76    init(connection: NWConnection, continuation: CheckedContinuation<Bool, Never>) {
77        self.connection = connection
78        self.continuation = continuation
79    }
80
81    nonisolated func finish(open: Bool) {
82        lock.lock()
83        guard !resumed else {
84            lock.unlock()
85            return
86        }
87        resumed = true
88        lock.unlock()
89
90        connection.cancel()
91        continuation.resume(returning: open)
92    }
93}