//
//  ARPPacket.swift
//  Tun
//
//  Created by 安礼成 on 2024/8/25.
//
import Foundation

struct ARPPacket {
    // ARP操作码
    enum Opcode: UInt16 {
        case request = 0x01
        case response = 0x02
        
        func isRequest() -> Bool {
            return self == .request
        }
        
        func isResponse() -> Bool {
            return self == .response
        }
    }
     
    var hardwareType: UInt16
    var protocolType: UInt16
    var hardwareSize: UInt8
    var protocolSize: UInt8
    var opcode: Opcode
    var senderMAC: Data
    var senderIP: UInt32
    var targetMAC: Data
    var targetIP: UInt32
    
    init(hardwareType: UInt16, protocolType: UInt16, hardwareSize: UInt8, protocolSize: UInt8, opcode: Opcode, 
         senderMAC: Data, senderIP: UInt32, targetMAC: Data, targetIP: UInt32) {
        
        self.hardwareType = hardwareType
        self.protocolType = protocolType
        self.hardwareSize = hardwareSize
        self.protocolSize = protocolSize
        self.opcode = opcode
        self.senderMAC = senderMAC
        self.senderIP = senderIP
        self.targetMAC = targetMAC
        self.targetIP = targetIP
    }

    init?(data: Data) {
        guard data.count >= 28 else {
            NSLog("length < 28: len: \(data.count)")
            return nil
        }
        
        self.hardwareType = UInt16(data[0]) << 8 | UInt16(data[1])
        self.protocolType = UInt16(data[2]) << 8 | UInt16(data[3])
        self.hardwareSize = data[4]
        self.protocolSize = data[5]
        guard let opcode = Opcode(rawValue: UInt16(data[6]) << 8 | UInt16(data[7])) else {
            NSLog("opcode error")
            return nil
        }
        
        self.opcode = opcode
        self.senderMAC = Data(data[8..<14])
        self.senderIP = UInt32(data: Data(data[14..<18]))
        self.targetMAC = Data(data[18..<24])
        self.targetIP = UInt32(data: Data(data[24..<28]))
    }
    
    func marshal() -> Data {
        var data = Data()
        data.append(self.hardwareType.data())
        data.append(self.protocolType.data())
        data.append(self.hardwareSize)
        data.append(self.protocolSize)
        data.append(self.opcode.rawValue.data())
        data.append(self.senderMAC)
        data.append(Data(uint32: self.senderIP))
        data.append(self.targetMAC)
        data.append(Data(uint32: self.targetIP))
        
        return data
    }
    
}

extension ARPPacket {
    
    static func arpRequest(senderIP: UInt32, senderMAC: Data, targetIP: UInt32) -> ARPPacket {
        return ARPPacket(hardwareType: 0x01,
                         protocolType: 0x0800,
                         hardwareSize: 0x06,
                         protocolSize: 0x04,
                         opcode: .request,
                         senderMAC: senderMAC,
                         senderIP: senderIP,
                         targetMAC: Data([0x00, 0x00, 0x00, 0x00, 0x00, 0x00]),
                         targetIP: targetIP)
    }
    
    static func arpResponse(for arp: ARPPacket, mac: Data, ip: UInt32) -> ARPPacket {
        return ARPPacket(hardwareType: arp.hardwareType,
                         protocolType: arp.protocolType,
                         hardwareSize: arp.hardwareSize,
                         protocolSize: arp.protocolSize,
                         opcode: .response,
                         senderMAC: mac,
                         senderIP: ip,
                         targetMAC: arp.senderMAC,
                         targetIP: arp.senderIP)
    }
        
}