//
//  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
    let transportPacket: TransportPacket
    
    enum TransportPacket {
        case tcp(TCPPacket)
        case udp(UDPPacket)
        case icmp(ICMPPacket)
        case unsupported(UInt8)
        case malformed
    }
    
    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
        self.transportPacket = Self.parseTransportPacket(proto: data[9], offset: Int(headerLen), data: data)
    }

    private static func parseTransportPacket(proto: UInt8, offset: Int, data: Data) -> TransportPacket {
        guard let proto = TransportProtocol(rawValue: proto) else {
            return .unsupported(proto)
        }
        
        switch proto {
        case .tcp:
            guard let tcp = TCPPacket(data, offset: offset) else {
                return .malformed
            }
            return .tcp(tcp)
            
        case .udp:
            guard let udp = UDPPacket(data, offset: offset) else {
                return .malformed
            }
            return .udp(udp)
            
        case .icmp:
            guard let icmp = ICMPPacket(data, offset: offset) else {
                return .malformed
            }
            return .icmp(icmp)
        }
    }
    
}

// MARK: - Lightweight IP Packet View

struct IPPacketView {
    let header: IPHeader
    let data: Data
    let transportPacket: TransportPacket

    enum TransportPacket {
        case tcp(srcPort: UInt16, dstPort: UInt16)
        case udp(srcPort: UInt16, dstPort: UInt16, payloadOffset: Int)
        case icmp
        case unsupported(UInt8)
        case malformed
    }

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

        let firstByte = Self.byte(at: 0, in: data)
        let version = firstByte >> 4
        let headerLen = (firstByte & 0x0F) * 4

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

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

        self.data = data
        self.transportPacket = Self.parseTransportPacket(proto: proto, offset: Int(headerLen), data: data)
    }

    private static func parseTransportPacket(proto: UInt8, offset: Int, data: Data) -> TransportPacket {
        guard let proto = TransportProtocol(rawValue: proto) else {
            return .unsupported(proto)
        }

        switch proto {
        case .tcp:
            guard data.count >= offset + 20 else {
                return .malformed
            }
            return .tcp(
                srcPort: UInt16(bytes: (Self.byte(at: offset, in: data), Self.byte(at: offset + 1, in: data))),
                dstPort: UInt16(bytes: (Self.byte(at: offset + 2, in: data), Self.byte(at: offset + 3, in: data)))
            )

        case .udp:
            guard data.count >= offset + 8 else {
                return .malformed
            }
            return .udp(
                srcPort: UInt16(bytes: (Self.byte(at: offset, in: data), Self.byte(at: offset + 1, in: data))),
                dstPort: UInt16(bytes: (Self.byte(at: offset + 2, in: data), Self.byte(at: offset + 3, in: data))),
                payloadOffset: offset + 8
            )

        case .icmp:
            guard data.count >= offset + 4 else {
                return .malformed
            }
            return .icmp
        }
    }

    private static func byte(at offset: Int, in data: Data) -> UInt8 {
        data[data.index(data.startIndex, offsetBy: offset)]
    }
}

// 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

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

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

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

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

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

        let window = UInt16(bytes: (data[offset + 14], data[offset + 15]))
        let checksum = UInt16(bytes: (data[offset + 16], data[offset + 17]))
        let urgent = UInt16(bytes: (data[offset + 18], data[offset + 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 >= offset + headerLen else {
            return nil
        }
        
        self.header = header
    }
    
}

// MARK: - UDP Packet

struct UDPPacket {
    let srcPort: UInt16
    let dstPort: UInt16
    let length: UInt16
    let checksum: UInt16
    let payloadOffset: Int
    
    init?(_ data: Data, offset: Int) {
        guard data.count >= offset + 8 else {
            return nil
        }
        
        self.srcPort = UInt16(bytes: (data[offset], data[offset + 1]))
        self.dstPort = UInt16(bytes: (data[offset + 2], data[offset + 3]))
        self.length  = UInt16(bytes: (data[offset + 4], data[offset + 5]))
        self.checksum = UInt16(bytes: (data[offset + 6], data[offset + 7]))
        self.payloadOffset = offset + 8
    }
    
}

// MARK: - ICMP Packet

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

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

        self.type = data[offset]
        self.code = data[offset + 1]
        self.checksum = UInt16(bytes: (data[offset + 2], data[offset + 3]))
    }
    
}