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

import Foundation

struct IPHeader {
    let version: UInt8
    let headerLength: UInt8
    let typeOfService: UInt8
    let totalLength: UInt16
    let id: UInt16
    let offset: UInt16
    let timeToLive: UInt8
    let proto: UInt8
    let checksum: UInt16
    let source: UInt32
    let destination: UInt32
    
    var source_ip: String {
        return SDLUtil.int32ToIp(source)
    }
    
    var destination_ip: String {
        return SDLUtil.int32ToIp(destination)
    }
    
    public var description: String {
        """
        IPHeader version: \(version), header length: \(headerLength), type of service: \(typeOfService), total length: \(totalLength),
        id: \(id), offset: \(offset), time ot live: \(timeToLive), proto: \(proto), checksum: \(checksum), source ip: \(source_ip), destination ip:\(destination_ip)
        """
    }
}

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

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

struct IPPacket {
    let header: IPHeader
    let data: Data
    
    init?(_ data: Data) {
        guard data.count >= 20 else {
            return nil
        }
        
        self.header = IPHeader(version: data[0] >> 4,
                               headerLength: (data[0] & 0b1111) * 4,
                               typeOfService: data[1],
                               totalLength: UInt16(bytes: (data[2], data[3])),
                               id: UInt16(bytes: (data[4], data[5])),
                               offset: 1,
                               timeToLive: 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
    }
    
    // 获取负载部分
    func getPayload() -> Data {
        return data.subdata(in: 20..<data.count)
    }
    
    // 获取ip数据包里面目标端口
    func getDstPort() -> UInt16? {
        guard case .ipv4 = IPVersion(rawValue: self.header.version), self.data.count >= 24 else {
            return nil
        }
        
        // 系统只会读取到ipv4的数据包，(srcPort:16, dstPort:16, ...)
        return UInt16(bytes: (self.data[22], self.data[23]))
    }
    
}
