ebpf-out-redirect.py

#!/usr/bin/python3
import argparse
from bcc import BPF
import pyroute2

def main():
    parser = argparse.ArgumentParser(description='Redirect traffic to another interface with VLAN')
    subparsers = parser.add_subparsers(dest='command')
    load_parser = subparsers.add_parser('load')
    load_parser.add_argument('--iface-in', required=True, help='Name of interface to attach BPF program')
    load_parser.add_argument('--iface-out', required=True, help='Name of interface to send traffic')
    load_parser.add_argument('--dst-mac', required=True, help='Destination MAC address')
    load_parser.add_argument('--src-mac', required=True, help='Source MAC address')
    load_parser.add_argument('--vlan', required=True, type=int, help='VLAN ID')
    unload_parser = subparsers.add_parser('unload')
    unload_parser.add_argument('--iface-in', required=True, help='Name of interface to unload BPF program')
    args = parser.parse_args()

    if args.command == 'load':
        load(args)
    elif args.command == 'unload':
        unload(args)
    else:
        print("Invalid command. Please use 'load' or 'unload'.")
        exit(1)

def load(args):
    bpf_src = """
    #include <linux/bpf.h>
    #include <linux/if_ether.h>
    #include <linux/ip.h>
    #include <linux/ipv6.h>
    #include <linux/pkt_cls.h>
    #include <bcc/proto.h>

    #define IF_INDEX $$IF_INDEX$$
    #define DST_MAC $$DST_MAC$$
    #define SRC_MAC $$SRC_MAC$$
    #define VLAN_ID $$VLAN_ID$$

    // ==== UNCOMMENT THIS TO TURN ON DEBUG ====
    // #define DEBUG 1

    #ifdef DEBUG
    #define bpf_debug(fmt, ...)                                                    \
        ({                                                                         \
            char ____fmt[] = fmt;                                                  \
            bpf_trace_printk(____fmt, sizeof(____fmt), ##__VA_ARGS__);             \
        })
    #else
    #define bpf_debug(fmt, ...)                                                    \
        {}                                                                         \
        while (0)
    #endif

    __always_inline static int is_global_unicast_ipv4(__u32 a) {
        __u32 b = a >> 24U;
        if (b < 0xe0) {
            if (b == 0x00) /* 0.0.0.0/8        This network */
                return 0;
            if (b == 0x7f) /* 127.0.0.0/8      Loopback address */
                return 0;
            if ((b == 0x0a) || /* 10.0.0.0/8       Private range */
                ((a & 0xffff0000) ==
                0xc0a80000) || /* 192.168.0.0/16   Private range */
                ((a & 0xfff00000) ==
                0xac100000)) /* 172.16.0.0/12    Private range */
                return 0;
            return 1;
        }
        if (b < 0xf0) /* 224.0.0.0/4      Multicast address */
            return 0;
        if (a == 0xffffffff) /* 255.255.255.255  Broadcast address */
            return 0;
        return 0; /* 240.0.0.0/4      Reserved / private */
    }

    __always_inline static int is_global_unicast_ipv6(__u8 *a) {
        // 2000::/3
        return (a[0] & 0xe0) == 0x20;
    }

    struct vlan_hdr {
        unsigned char h_dest[6];
        unsigned char h_source[6];
        __be16 h_vlan_proto;
        __be16 h_vlan_TCI;
    } __attribute__((packed));

    int tc_egress(struct __sk_buff *skb) {
        void *data = (void *)(long)skb->data;
        void *data_end = (void *)(long)skb->data_end;

        if (skb->protocol == bpf_htons(ETH_P_IP)) {
            if (data + sizeof(struct ethhdr) + sizeof(struct iphdr) > data_end) {
                bpf_debug("Skip malformed IPv4 packet");
                return TC_ACT_OK;
            }
            struct iphdr *ipv4 = (struct iphdr *)((struct ethhdr *)data + 1);
            __u32 daddr = bpf_ntohl(ipv4->daddr);
            if (!is_global_unicast_ipv4(daddr)) {
                bpf_debug("Skip non global unicast IPv4 packet to %llu",
                        (((__u64)daddr >> 24) & 0xff) * 1000000000 +
                            (((__u64)daddr >> 16) & 0xff) * 1000000 +
                            (((__u64)daddr >> 8) & 0xff) * 1000 +
                            (((__u64)daddr >> 0) & 0xff));
                return TC_ACT_OK;
            }

        } else if (skb->protocol == bpf_htons(ETH_P_IPV6)) {
            if (data + sizeof(struct ethhdr) + sizeof(struct ipv6hdr) > data_end) {
                return TC_ACT_OK;
            }
            struct ipv6hdr *ipv6 = (struct ipv6hdr *)((struct ethhdr *)data + 1);
            if (!is_global_unicast_ipv6(ipv6->daddr.s6_addr)) {
                bpf_debug("Skip non global unicast IPv6 packet to %llx:%llx",
                        ((__u64)bpf_ntohl(ipv6->daddr.s6_addr32[0]) << 32) |
                            bpf_ntohl(ipv6->daddr.s6_addr32[1]),
                        ((__u64)bpf_ntohl(ipv6->daddr.s6_addr32[2]) << 32) |
                            bpf_ntohl(ipv6->daddr.s6_addr32[3]));
                return TC_ACT_OK;
            }
        } else {
            // Skip non IPv4/IPv6 traffic
            return TC_ACT_OK;
        }

        // Change dst/src MAC and assign VLAN
        bpf_skb_change_head(skb, 4, 0);

        struct vlan_hdr new_hdr = {
            .h_dest = {DST_MAC},
            .h_source = {SRC_MAC},
            .h_vlan_proto = bpf_htons(ETH_P_8021Q),
            .h_vlan_TCI = bpf_htons(VLAN_ID),
        };

        if (skb->data + sizeof(new_hdr) > skb->data_end) {
            return TC_ACT_SHOT;
        }

        bpf_skb_store_bytes(skb, 0, &new_hdr, sizeof(new_hdr), 0);

        return bpf_redirect(IF_INDEX, 0);
    }
    """

    ipr = pyroute2.IPRoute()

    mac_format = lambda mac: "0x" + ', 0x'.join(mac.split(':'))
    bpf_src = bpf_src.replace("$$IF_INDEX$$", str(ipr.link_lookup(ifname=args.iface_out)[0]))
    bpf_src = bpf_src.replace("$$DST_MAC$$", mac_format(args.dst_mac))
    bpf_src = bpf_src.replace("$$SRC_MAC$$", mac_format(args.src_mac))
    bpf_src = bpf_src.replace("$$VLAN_ID$$", str(args.vlan))

    prog = BPF(text=bpf_src)
    tc_egress_fn = prog.load_func("tc_egress", BPF.SCHED_CLS)

    if_index = ipr.link_lookup(ifname=args.iface_in)[0]
    ipr.tc("add", "clsact", if_index)
    ipr.tc("add-filter", "bpf", if_index, ":1", fd=tc_egress_fn.fd, name=tc_egress_fn.name, parent="ffff:fff3", classid=1, direct_action=True)

def unload(args):
    ipr = pyroute2.IPRoute()
    input_idx = ipr.link_lookup(ifname=args.iface_in)[0]
    ipr.tc("del", "clsact", input_idx)

if __name__ == "__main__":
    main()