package parser

import (
    "encoding/binary"
    "errors"
    "fmt"
    "net"
)

const (
    IPv4ProtoICMP  = 1   //互联网控制消息协议
    IPv4ProtoIGMP  = 2   //互联网组管理协议
    IPv4ProtoTCP   = 6   //传输控制协议
    IPv4ProtoUDP   = 17  //用户数据报协议
    IPv4ProtoENCAP = 41  //IPv6封装
    IPv4ProtoOSPF  = 89  //开放式最短路径优先
    IPv4ProtoSCTP  = 132 //流控制传输协议
)

type IPv4 struct {
    BaseLayer
    Version     uint8        // 版本
    IHL         uint8        // 首部长度
    DS          uint8        // 区分服务
    TotalLength uint16       // 总长
    Id          uint16       // 标识符
    Flags       uint8        // 标志
    FragOffset  uint16       // 分片偏移
    TTL         uint8        // 存活时间
    Protocol    uint8        // 协议
    Checksum    uint16       // 首部校验和
    SrcIP       net.IP       // 源地址
    DstIP       net.IP       // 目的地址
    Options     []IPv4Option // 选项字段
    Padding     []byte       // 填充
}

type IPv4Option struct {
    OptionType   uint8
    OptionLength uint8
    OptionData   []byte
}

func (I *IPv4) LayerType() int {
    return LayerTypeIPv4
}

// DecodeIPV4 ref: https://zh.wikipedia.org/wiki/IPv4#%E6%8A%A5%E6%96%87%E7%BB%93%E6%9E%84
func DecodeIPV4(data []byte) (*IPv4, error) {
    // 至少要有20bytes
    if len(data) < 20 {
        return nil, fmt.Errorf("[ipv4] 无效的长度：%d, 至少需要20", len(data))
    }

    ipv4 := IPv4{}
    ipv4.Version = data[0] >> 4
    ipv4.IHL = uint8(data[0]) & 0x0F
    ipv4.DS = data[1]
    ipv4.TotalLength = binary.BigEndian.Uint16(data[2:4])
    ipv4.Id = binary.BigEndian.Uint16(data[4:6])
    flagsFrags := binary.BigEndian.Uint16(data[6:8])
    ipv4.Flags = uint8(flagsFrags >> 13)
    ipv4.FragOffset = flagsFrags & 0x1FFF
    ipv4.TTL = data[8]
    ipv4.Protocol = data[9]
    ipv4.Checksum = binary.BigEndian.Uint16(data[10:12])
    ipv4.SrcIP = data[12:16]
    ipv4.DstIP = data[16:20]

    if ipv4.TotalLength == 0 {
        ipv4.TotalLength = uint16(len(data))
    }

    // 检查数据
    if ipv4.TotalLength < 20 {
        return nil, fmt.Errorf("[IPv4] 数据包长度太小 (%d < 20)", ipv4.TotalLength)
    } else if ipv4.IHL < 5 {
        return nil, fmt.Errorf("[IPv4] 标头长度太小 (%d < 5)", ipv4.IHL)
    } else if int(ipv4.IHL*4) > int(ipv4.TotalLength) {
        return nil, fmt.Errorf("[IPv4] 标头长度 > IP 总长度 (%d > %d)", ipv4.IHL, ipv4.TotalLength)
    }

    // 检查数据实际长度是否大于声明长度，大于（有padding）则除去padding拿到实际长度的payload
    if cmp := len(data) - int(ipv4.TotalLength); cmp > 0 {
        data = data[:ipv4.TotalLength]
    } else if cmp < 0 {
        // 如果声明的头部长度比整体数据还大，异常
        if int(ipv4.IHL)*4 > len(data) {
            return nil, errors.New("[ipv4] 数据包标头长度大于总数据包长度，标头长度无效或某些数据丢失")
        }
    }

    ipv4.Header = data[:ipv4.IHL*4]
    ipv4.Payload = data[ipv4.IHL*4:]

    // 20bytes固定头部往后到IHL都是选项字段了
    data = data[20 : ipv4.IHL*4]

    // 解析可选字段，如果data[20 : ipv4.IHL*4]之间还有数据，说明选项字段有内容
    for len(data) > 0 {
        if ipv4.Options == nil {
            // Pre-allocate to avoid growing the slice too much.
            ipv4.Options = make([]IPv4Option, 0, 4)
        }
        opt := IPv4Option{OptionType: data[0]}
        switch opt.OptionType {
        case 0: // option字段结束，EOL标志
            opt.OptionLength = 1
            ipv4.Options = append(ipv4.Options, opt)
            ipv4.Padding = data[1:]
            goto FINISH
        case 1: // 1byte填充
            opt.OptionLength = 1
            data = data[1:]
            ipv4.Options = append(ipv4.Options, opt)
        default:
            // 字段未结束且不是1byte填充字段长度不足2byte
            if len(data) < 2 {
                return nil, fmt.Errorf("[IPv4] 选项长度 %d 小于2", len(data))
            }
            opt.OptionLength = data[1]
            if len(data) < int(opt.OptionLength) {
                return nil, fmt.Errorf("[IPv4] 可选字段长度超过剩余 IP 标头大小，选项类型:%v, 长度:%v", opt.OptionType, opt.OptionLength)
            }
            if opt.OptionLength <= 2 {
                return nil, fmt.Errorf("[IPv4] 可选字段长度小于2，选项类型:%v 长度:%d", opt.OptionType, opt.OptionLength)
            }
            opt.OptionData = data[2:opt.OptionLength]
            data = data[opt.OptionLength:]
            ipv4.Options = append(ipv4.Options, opt)
        }
    }

FINISH:
    return &ipv4, nil
}