Friday, December 30, 2011

Raw sockets with BPF in Python

Update 2021: Note that this was written in 2011. Nowadays I'd not recommend doing it this way, but instead using BCC to write your filters in C, from within Python. The following example shows the use of raw sockets where filtering is applied using BPF. It has only been tested on Linux. The filter data structure is built up to form a machine language-like set of instructions that decide whether the packet should be passed to the raw socket or not. The filter is applied to the socket using SO_ATTACH_FILTER. This specific example filters packets to only allow packets destined a DHCP server (UDP port 67).
from binascii import hexlify
from ctypes import create_string_buffer, addressof
from socket import socket, AF_PACKET, SOCK_RAW, SOL_SOCKET
from struct import pack, unpack


# A subset of Berkeley Packet Filter constants and macros, as defined in
# linux/filter.h.

# Instruction classes
BPF_LD = 0x00
BPF_JMP = 0x05
BPF_RET = 0x06

# ld/ldx fields
BPF_H = 0x08
BPF_B = 0x10
BPF_ABS = 0x20

# alu/jmp fields
BPF_JEQ = 0x10
BPF_K = 0x00

def bpf_jump(code, k, jt, jf):
    return pack('HBBI', code, jt, jf, k)

def bpf_stmt(code, k):
    return bpf_jump(code, k, 0, 0)


# Ordering of the filters is backwards of what would be intuitive for 
# performance reasons: the check that is most likely to fail is first.
filters_list = [
    # Must have dst port 67. Load (BPF_LD) a half word value (BPF_H) in 
    # ethernet frame at absolute byte offset 36 (BPF_ABS). If value is equal to
    # 67 then do not jump, else jump 5 statements.
    bpf_stmt(BPF_LD | BPF_H | BPF_ABS, 36),
    bpf_jump(BPF_JMP | BPF_JEQ | BPF_K, 67, 0, 5),

    # Must be UDP (check protocol field at byte offset 23)
    bpf_stmt(BPF_LD | BPF_B | BPF_ABS, 23), 
    bpf_jump(BPF_JMP | BPF_JEQ | BPF_K, 0x11, 0, 3),

    # Must be IPv4 (check ethertype field at byte offset 12)
    bpf_stmt(BPF_LD | BPF_H | BPF_ABS, 12), 
    bpf_jump(BPF_JMP | BPF_JEQ | BPF_K, 0x0800, 0, 1),

    bpf_stmt(BPF_RET | BPF_K, 0x0fffffff), # pass
    bpf_stmt(BPF_RET | BPF_K, 0), # reject
]

# Create filters struct and fprog struct to be used by SO_ATTACH_FILTER, as
# defined in linux/filter.h.
filters = ''.join(filters_list)
b = create_string_buffer(filters)
mem_addr_of_filters = addressof(b)
fprog = pack('HL', len(filters_list), mem_addr_of_filters)

# As defined in asm/socket.h
SO_ATTACH_FILTER = 26

# Create listening socket with filters
s = socket(AF_PACKET, SOCK_RAW, 0x0800)
s.setsockopt(SOL_SOCKET, SO_ATTACH_FILTER, fprog)
s.bind(('eth0', 0x0800))

while True:
    data, addr = s.recvfrom(65565)
    print 'got data from', addr, ':', hexlify(data)

Thursday, December 08, 2011

Simple heap implementation in C#

I recently needed a simple priority queue in C# and was surprised to find that there is neither anything called "priority queue" nor "min heap" or "max heap" in .NET (or at least I did not find it when I looked). Here is a textbook-like implementation I came up with:

class MinHeap<T> where T : IComparable
{
    private List<T> data = new List<T>();

    public void Insert(T o)
    {
        data.Add(o);

        int i = data.Count - 1;
        while (i > 0)
        {
            int j = (i + 1) / 2 - 1;

            // Check if the invariant holds for the element in data[i]
            T v = data[j];
            if (v.CompareTo(data[i]) < 0 || v.CompareTo(data[i]) == 0)
            {
                break;
            }

            // Swap the elements
            T tmp = data[i];
            data[i] = data[j];
            data[j] = tmp;

            i = j;
        }
    }

    public T ExtractMin()
    {
        if (data.Count < 0)
        {
            throw new ArgumentOutOfRangeException();
        }

        T min = data[0];
        data[0] = data[data.Count - 1];
        data.RemoveAt(data.Count - 1);
        this.MinHeapify(0);
        return min;
    }

    public T Peek()
    {
        return data[0];
    }

    public int Count
    {
        get { return data.Count; }
    }

    private void MinHeapify(int i)
    {
        int smallest;
        int l = 2 * (i + 1) - 1;
        int r = 2 * (i + 1) - 1 + 1;

        if (l < data.Count && (data[l].CompareTo(data[i]) < 0))
        {
            smallest = l;
        }
        else
        {
            smallest = i;
        }

        if (r < data.Count && (data[r].CompareTo(data[smallest]) < 0))
        {
            smallest = r;
        }

        if (smallest != i)
        {
            T tmp = data[i];
            data[i] = data[smallest];
            data[smallest] = tmp;
            this.MinHeapify(smallest);
        }
    }
}