Add Wireshark packet dissector plugin and associated tools

scripts/packet-analysis/dev/PcapTools/src/PcapTools.jl

@@ -0,0 +1,26 @@
+module PcapTools
+
+using Dates
+using Mmap
+using UnixTimes
+using UnsafeArrays
+
+export PcapHeader, RecordHeader
+export PcapRecord
+export PcapReader, PcapStreamReader, PcapBufferReader
+export PcapWriter, PcapStreamWriter
+export LINKTYPE_NULL, LINKTYPE_ETHERNET
+export splitcap
+
+abstract type PcapReader end
+abstract type PcapWriter end
+
+include("pcap_header.jl")
+include("record_header.jl")
+include("record.jl")
+include("buffer_reader.jl")
+include("stream_reader.jl")
+include("stream_writer.jl")
+include("splitcap.jl")
+
+end

scripts/packet-analysis/dev/PcapTools/src/buffer_reader.jl

@@ -0,0 +1,96 @@
+"""
+Reads pcap data from an array of bytes.
+"""
+mutable struct PcapBufferReader <: PcapReader
+    data::Vector{UInt8}
+    raw_header::Vector{UInt8}
+    header::PcapHeader
+    offset::Int64
+    mark::Int64
+    usec_mul::Int64
+    bswapped::Bool
+
+    @doc """
+        PcapBufferReader(data::Vector{UInt8})
+
+    Create reader over `data`. Will read and process pcap header,
+    and yield records through `read(::PcapBufferReader)`.
+    """
+    function PcapBufferReader(data::Vector{UInt8})
+        length(data) < sizeof(PcapHeader) && throw(EOFError())
+        rh = data[1:sizeof(PcapHeader)]
+        h = unsafe_load(Ptr{PcapHeader}(pointer(data)))
+        h, bswapped, nanotime = process_header(h)
+        new(data, rh, h, sizeof(h), -1, nanotime ? 1 : 1000, bswapped)
+    end
+end
+
+"""
+    PcapBufferReader(path::AbstractString)
+
+Memory map file in `path` and create PcapBufferReader over its content.
+"""
+function PcapBufferReader(path::AbstractString)
+    io = open(path)
+    data = Mmap.mmap(io)
+    PcapBufferReader(data)
+end
+
+function Base.close(x::PcapBufferReader)
+    x.data = UInt8[]
+    x.offset = 0
+    nothing
+end
+
+Base.length(x::PcapBufferReader) = length(x.data)
+Base.position(x::PcapBufferReader) = x.offset; nothing
+Base.seek(x::PcapBufferReader, pos) = x.offset = pos; nothing
+
+function Base.mark(x::PcapBufferReader)
+    x.mark = x.offset
+    x.mark
+end
+
+function Base.unmark(x::PcapBufferReader)
+    if x.mark >= 0
+        x.mark = -1
+        true
+    else
+        false
+    end
+end
+
+Base.ismarked(x::PcapBufferReader) = x.mark >= 0
+
+function Base.reset(x::PcapBufferReader)
+    !ismarked(x) && error("PcapBufferReader not marked")
+    x.offset = x.mark
+    x.mark = -1
+    x.offset
+end
+
+Base.eof(x::PcapBufferReader) = (length(x) - x.offset) < sizeof(RecordHeader)
+
+"""
+    read(x::PcapBufferReader) -> PcapRecord
+
+Read one record from pcap data.
+Throws `EOFError` if no more data available.
+"""
+@inline function Base.read(x::PcapBufferReader)
+    eof(x) && throw(EOFError())
+    record_offset = x.offset
+    p = pointer(x.data) + record_offset
+    GC.@preserve x begin
+        h = unsafe_load(Ptr{RecordHeader}(p))
+    end
+    if x.bswapped
+        h = bswap(h)
+    end
+    t1 = (h.ts_sec + x.header.thiszone) * 1_000_000_000
+    t2 = Int64(h.ts_usec) * x.usec_mul
+    t = UnixTime(Dates.UTInstant(Nanosecond(t1 + t2)))
+    x.offset += sizeof(RecordHeader) + h.incl_len
+    x.offset > length(x) && error("Insufficient data in pcap record")
+    PcapRecord(h, t, x.data, record_offset)
+end

scripts/packet-analysis/dev/PcapTools/src/pcap_header.jl

@@ -0,0 +1,46 @@
+# NOTE: Not using @enum because it craps out when displaying unknown values
+const LINKTYPE_NULL = UInt32(0)
+const LINKTYPE_ETHERNET = UInt32(1)
+
+struct PcapHeader
+    magic::UInt32
+    version_major::UInt16
+    version_minor::UInt16
+    thiszone::Int32
+    sigfigs::UInt32
+    snaplen::UInt32
+    linktype::UInt32
+end
+
+function Base.bswap(x::PcapHeader)
+    PcapHeader(
+        bswap(x.magic),
+        bswap(x.version_major),
+        bswap(x.version_minor),
+        bswap(x.thiszone),
+        bswap(x.sigfigs),
+        bswap(x.snaplen),
+        bswap(x.linktype))
+end
+
+function process_header(x::PcapHeader)
+    if x.magic == 0xa1b2c3d4
+        bswapped = false
+        nanotime = false
+    elseif x.magic == 0xd4c3b2a1
+        bswapped = true
+        nanotime = false
+    elseif x.magic == 0xa1b23c4d
+        bswapped = false
+        nanotime = true
+    elseif x.magic == 0x4d3cb2a1
+        bswapped = true
+        nanotime = true
+    else
+        throw(ArgumentError("Invalid pcap header"))
+    end
+    if bswapped
+        x = bswap(x)
+    end
+    x, bswapped, nanotime
+end

scripts/packet-analysis/dev/PcapTools/src/record.jl

@@ -0,0 +1,25 @@
+"""
+Record of pcap data
+"""
+struct PcapRecord
+    header::RecordHeader
+    timestamp::UnixTime
+    underlying_data::Vector{UInt8}
+    record_offset::Int
+end
+
+@inline function record_field_(x::PcapRecord, ::Val{:data})
+    offset = getfield(x, :record_offset) + sizeof(RecordHeader)
+    len = Int(getfield(x, :header).incl_len)
+    UnsafeArray{UInt8, 1}(pointer(getfield(x, :underlying_data)) + offset, (len,))
+end
+
+@inline function record_field_(x::PcapRecord, ::Val{:raw})
+    offset = getfield(x, :record_offset)
+    len = sizeof(RecordHeader) + getfield(x, :header).incl_len
+    UnsafeArray{UInt8, 1}(pointer(getfield(x, :underlying_data)) + offset, (len,))
+end
+
+@inline record_field_(x::PcapRecord, ::Val{f}) where {f} = getfield(x, f) 
+
+@inline Base.getproperty(x::PcapRecord, f::Symbol) = record_field_(x, Val(f))

scripts/packet-analysis/dev/PcapTools/src/record_header.jl

@@ -0,0 +1,14 @@
+struct RecordHeader
+    ts_sec::UInt32
+    ts_usec::UInt32
+    incl_len::UInt32
+    orig_len::UInt32
+end
+
+function Base.bswap(x::RecordHeader)
+    RecordHeader(
+        bswap(x.ts_sec),
+        bswap(x.ts_usec),
+        bswap(x.incl_len),
+        bswap(x.orig_len))
+end

scripts/packet-analysis/dev/PcapTools/src/splitcap.jl

@@ -0,0 +1,102 @@
+strip_nothing_(::Type{Union{Nothing, T}}) where T = T
+strip_nothing_(::Type{T}) where T = T
+
+progress_noop_(n) = nothing
+
+mutable struct SplitCapOutput{S}
+    work_buffer::Vector{UInt8}
+    complete_buffers::Channel{Vector{UInt8}}
+    stream::S
+end
+
+# Since Julia (as of 1.5) doesn't support task migration,
+# continue with new task after each buffer write, to rebalance across threads
+function write_one_and_continue_(output::SplitCapOutput, free_buffers::Channel, pending::Threads.Atomic{Int})
+    i = iterate(output.complete_buffers)
+    if i === nothing
+        Threads.atomic_sub!(pending, 1)
+        return nothing
+    end
+    b, _ = i
+    write(output.stream, b)
+    empty!(b)
+    put!(free_buffers, b)
+    Threads.@spawn write_one_and_continue_($output, $free_buffers, $pending)
+end
+
+function splitcap(
+    ::Type{KeyType},
+    ::Type{StreamType},
+    reader::PcapReader,
+    record2key,
+    key2stream,
+    progress_callback = progress_noop_;
+    own_streams::Bool = true
+) where {KeyType, StreamType}
+    buffer_size = 1024 * 1024 * 2
+    max_pending_buffers = 4
+    outputs = Dict{KeyType, SplitCapOutput{StreamType}}()
+    free_buffers = Channel{Vector{UInt8}}(Inf)
+    n = 0
+    pending = Threads.Atomic{Int}(0)
+    try
+        while !eof(reader)
+            record = read(reader)
+            dst = record2key(record)
+            if dst isa KeyType
+                output = get!(outputs, dst) do
+                    stream = key2stream(dst)
+                    buffer = sizehint!(UInt8[], buffer_size + 1500)
+                    own_streams && append!(buffer, reader.raw_header)
+                    output = SplitCapOutput{StreamType}(
+                        buffer,
+                        Channel{Vector{UInt8}}(max_pending_buffers),
+                        stream)
+                    Threads.atomic_add!(pending, 1)
+                    Threads.@spawn write_one_and_continue_($output, $free_buffers, $pending)
+                    output
+                end
+                append!(output.work_buffer, record.raw)
+                if length(output.work_buffer) >= buffer_size
+                    put!(output.complete_buffers, output.work_buffer)
+                    if isready(free_buffers)
+                        output.work_buffer = take!(free_buffers)
+                    else
+                        output.work_buffer = sizehint!(UInt8[], buffer_size + 1500)
+                    end
+                end
+            end
+            n += 1
+            progress_callback(n)
+            GC.safepoint()
+        end
+        for output in values(outputs)
+            if !isempty(output.work_buffer)
+                put!(output.complete_buffers, output.work_buffer)
+            end
+            close(output.complete_buffers)
+        end
+        while pending[] != 0
+            sleep(0.1)
+        end
+    finally
+        if own_streams
+            for output in values(outputs)
+                close(output.stream)
+            end
+        end
+    end
+    nothing
+end
+
+function splitcap(
+    reader::PcapReader,
+    record2key,
+    key2stream,
+    progress_callback = progress_noop_;
+    kwargs...
+)
+    KeyType = strip_nothing_(Core.Compiler.return_type(record2key, Tuple{PcapRecord}))
+    StreamType = Core.Compiler.return_type(key2stream, Tuple{KeyType})
+    splitcap(KeyType, StreamType, reader, record2key, key2stream, progress_callback; kwargs...)
+end

scripts/packet-analysis/dev/PcapTools/src/stream_reader.jl

@@ -0,0 +1,63 @@
+"""
+Reads pcap data from a stream.
+"""
+mutable struct PcapStreamReader{Src <: IO} <: PcapReader
+    src::Src
+    raw_header::Vector{UInt8}
+    header::PcapHeader
+    usec_mul::Int64
+    bswapped::Bool
+    record_buffer::Vector{UInt8}
+
+    @doc """
+        PcapStreamReader(src::IO)
+
+    Create reader over `src`. Will read and process pcap header,
+    and yield records through `read(::PcapStreamReader)`.
+    """
+    function PcapStreamReader(src::Src) where {Src <: IO}
+        raw_header = read(src, sizeof(PcapHeader))
+        length(raw_header) != sizeof(PcapHeader) && throw(EOFError())
+        h = GC.@preserve raw_header unsafe_load(Ptr{PcapHeader}(pointer(raw_header)))
+        header, bswapped, nanotime = process_header(h)
+        new{Src}(src, raw_header, header, nanotime ? 1 : 1000, bswapped, zeros(UInt8, 9000 + sizeof(RecordHeader)))
+    end
+end
+
+"""
+    PcapStreamReader(path)
+
+Open file at `path` and create PcapStreamReader over its content.
+"""
+PcapStreamReader(path::AbstractString) = PcapStreamReader(open(path))
+
+Base.close(x::PcapStreamReader) = close(x.src)
+Base.position(x::PcapStreamReader) = position(x.src)
+Base.seek(x::PcapStreamReader, pos) = seek(x.src, pos)
+Base.mark(x::PcapStreamReader) = mark(x.src)
+Base.unmark(x::PcapStreamReader) = unmark(x.src)
+Base.ismarked(x::PcapStreamReader) = ismarked(x.src)
+Base.reset(x::PcapStreamReader) = reset(x.src)
+Base.eof(x::PcapStreamReader) = eof(x.src)
+
+"""
+    read(x::PcapStreamReader) -> PcapRecord
+
+Read one record from pcap data. Record is valid until next read().
+Throws `EOFError` if no more data available.
+"""
+function Base.read(x::PcapStreamReader)
+    p = pointer(x.record_buffer)
+    GC.@preserve x begin
+        unsafe_read(x.src, p, sizeof(RecordHeader))
+        h = unsafe_load(Ptr{RecordHeader}(p))
+        if x.bswapped
+            h = bswap(h)
+        end
+        unsafe_read(x.src, p + sizeof(RecordHeader), h.incl_len)
+    end
+    t1 = (h.ts_sec + x.header.thiszone) * 1_000_000_000
+    t2 = Int64(h.ts_usec) * x.usec_mul
+    t = UnixTime(Dates.UTInstant(Nanosecond(t1 + t2)))
+    PcapRecord(h, t, x.record_buffer, 0)
+end

scripts/packet-analysis/dev/PcapTools/src/stream_writer.jl

@@ -0,0 +1,30 @@
+struct PcapStreamWriter{Dst <: IO} <: PcapWriter
+    dst::Dst
+
+    function PcapStreamWriter{Dst}(dst::Dst; thiszone = 0, snaplen = 65535, linktype = LINKTYPE_ETHERNET) where {Dst <: IO}
+        h = PcapHeader(
+            0xa1b23c4d,
+            0x0002,
+            0x0004,
+            thiszone,
+            0,
+            snaplen,
+            linktype)
+            write(dst, reinterpret(UInt8, [h]))
+        new(dst)
+    end
+end
+
+PcapStreamWriter(io::IO; kwargs...) = PcapStreamWriter{typeof(io)}(io; kwargs...)
+PcapStreamWriter(path::AbstractString; kwargs...) = PcapStreamWriter(open(path, "w"); kwargs...)
+
+Base.close(x::PcapStreamWriter) = close(x.dst)
+
+function Base.write(x::PcapStreamWriter, timestamp::UnixTime, data)
+    sec, nsec = fldmod(Dates.value(timestamp), 1_000_000_000)
+    data_length = length(data)
+    h = RecordHeader(sec, nsec, data_length, data_length)
+    write(x.dst, reinterpret(UInt8, [h]))
+    write(x.dst, collect(data))
+    nothing
+end