punchnet-macos/Tun/Punchnet/NetworkStack/IPPacket.swift

//
//  IPPacket.swift
//  Tun
//
//  Created by 安礼成 on 2024/1/18.
//

import Foundation

enum IPVersion: UInt8 {
    case ipv4 = 4
    case ipv6 = 6
}

enum TransportProtocol: UInt8 {
    case icmp = 1
    case tcp  = 6
    case udp  = 17
}

// MARK: - IP Header

struct IPHeader {
    let version: UInt8
    let headerLength: UInt8
    let typeOfService: UInt8
    let totalLength: UInt16
    let id: UInt16
    let offset: UInt16
    let ttl: UInt8
    let proto: UInt8
    let checksum: UInt16
    let source: UInt32
    let destination: UInt32

    var headerBytes: Int {
        Int(headerLength)
    }
}

// MARK: - IP Packet

struct IPPacket {
    let header: IPHeader
    let data: Data
    
    enum TransportPacket {
        case tcp(TCPPacket)
        case udp(UDPPacket)
        case icmp(ICMPPacket)
        case unsupported(UInt8)
        case malformed
    }
    
    var transportPacket: TransportPacket {
        guard let proto = TransportProtocol(rawValue: header.proto) else {
            return .unsupported(header.proto)
        }
        
        switch proto {
        case .tcp:
            guard let tcp = TCPPacket(payload) else {
                return .malformed
            }
            return .tcp(tcp)
            
        case .udp:
            guard let udp = UDPPacket(payload) else {
                return .malformed
            }
            return .udp(udp)
            
        case .icmp:
            guard let icmp = ICMPPacket(payload) else {
                return .malformed
            }
            return .icmp(icmp)
        }
    }
    
    var payload: Data.SubSequence {
        let offset = Int(header.headerLength)
        
        return data[offset..<data.count]
    }

    init?(_ data: Data) {
        guard data.count >= 20 else {
            return nil
        }

        let version = data[0] >> 4
        let headerLen = (data[0] & 0x0F) * 4

        guard data.count >= headerLen else {
            return nil
        }

        self.header = IPHeader(
            version: version,
            headerLength: headerLen,
            typeOfService: data[1],
            totalLength: UInt16(bytes: (data[2], data[3])),
            id: UInt16(bytes: (data[4], data[5])),
            offset: UInt16(bytes: (data[6], data[7])),
            ttl: data[8],
            proto: data[9],
            checksum: UInt16(bytes: (data[10], data[11])),
            source: UInt32(bytes: (data[12], data[13], data[14], data[15])),
            destination: UInt32(bytes: (data[16], data[17], data[18], data[19]))
        )

        self.data = data
    }
}

// MARK: - TCP Flags

struct TCPFlags: OptionSet {
    let rawValue: UInt16

    static let fin = TCPFlags(rawValue: 1 << 0)
    static let syn = TCPFlags(rawValue: 1 << 1)
    static let rst = TCPFlags(rawValue: 1 << 2)
    static let psh = TCPFlags(rawValue: 1 << 3)
    static let ack = TCPFlags(rawValue: 1 << 4)
    static let urg = TCPFlags(rawValue: 1 << 5)
    static let ece = TCPFlags(rawValue: 1 << 6)
    static let cwr = TCPFlags(rawValue: 1 << 7)
}

// MARK: - TCP Header

struct TCPHeader {
    let srcPort: UInt16
    let dstPort: UInt16

    let seq: UInt32
    let ack: UInt32

    let dataOffset: UInt8
    let flags: TCPFlags

    let window: UInt16
    let checksum: UInt16
    let urgentPointer: UInt16

    var headerLength: Int {
        Int(dataOffset) * 4
    }
}

// MARK: - TCP Packet

struct TCPPacket {
    let header: TCPHeader
    let payload: Data

    init?(_ data: Data) {
        guard data.count >= 20 else {
            return nil
        }

        let srcPort = UInt16(bytes: (data[0], data[1]))
        let dstPort = UInt16(bytes: (data[2], data[3]))

        let seq = UInt32(bytes: (data[4], data[5], data[6], data[7]))
        let ack = UInt32(bytes: (data[8], data[9], data[10], data[11]))

        let offsetAndFlags = UInt16(bytes: (data[12], data[13]))

        let dataOffset = UInt8(offsetAndFlags >> 12)
        let flags = TCPFlags(rawValue: offsetAndFlags & 0x01FF)

        let window = UInt16(bytes: (data[14], data[15]))
        let checksum = UInt16(bytes: (data[16], data[17]))
        let urgent = UInt16(bytes: (data[18], data[19]))

        let header = TCPHeader(
            srcPort: srcPort,
            dstPort: dstPort,
            seq: seq,
            ack: ack,
            dataOffset: dataOffset,
            flags: flags,
            window: window,
            checksum: checksum,
            urgentPointer: urgent
        )

        let headerLen = header.headerLength

        guard data.count >= headerLen else {
            return nil
        }

        self.header = header
        self.payload = data[headerLen..<data.count]
    }
}

// MARK: - UDP Packet

struct UDPPacket {
    let srcPort: UInt16
    let dstPort: UInt16
    let length: UInt16
    let checksum: UInt16
    let payload: Data

    init?(_ data: Data) {

        guard data.count >= 8 else {
            return nil
        }

        self.srcPort = UInt16(bytes: (data[0], data[1]))
        self.dstPort = UInt16(bytes: (data[2], data[3]))
        self.length  = UInt16(bytes: (data[4], data[5]))
        self.checksum = UInt16(bytes: (data[6], data[7]))

        self.payload = data[8..<data.count]
    }
}

// MARK: - ICMP Packet

struct ICMPPacket {
    let type: UInt8
    let code: UInt8
    let checksum: UInt16
    let payload: Data

    init?(_ data: Data) {
        guard data.count >= 4 else {
            return nil
        }

        self.type = data[0]
        self.code = data[1]
        self.checksum = UInt16(bytes: (data[2], data[3]))
        self.payload = data[4..<data.count]
    }
}