/*
 * ipfw3x.c - raw IP_FW_X sockopt client for ipfw3 rule injection.
 *
 * Speaks the same ABI as /sbin/ipfw3 do_set_x()/do_get_x():
 *   setsockopt(IPPROTO_IP, IP_FW_X, {x_header{opcode}, payload}, 4+payload)
 *
 * Instruction encoding: every ipfw_insn is 8 bytes = 2 x 32-bit words
 * (LEN_OF_IPFWINSN == 2, sys/net/ipfw3/ip_fw3.h:85-86); F_LEN(cmd) is in
 * words including the instruction header word pair.
 *
 * usage:
 *   ipfw3x add <rulenum> <allow|deny|count|skipto> [arg1]
 *   ipfw3x add <rulenum> protodeny <proto>   # filter: proto N; action: deny
 *   ipfw3x del <rulenum>
 *   ipfw3x flush
 *   ipfw3x show
 *
 * Built for the DF-3439 skipto next_rule stale-cache UAF PoC.
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <errno.h>
#include <unistd.h>

#define IP_FW_X    49
#define IP_FW_ADD  50
#define IP_FW_DEL  51
#define IP_FW_FLUSH 52
#define IP_FW_GET  54

/* must byte-match sys/net/ipfw3/ip_fw3.h */
typedef struct _ipfw_insn {
	uint8_t  opcode;
	uint8_t  len;		/* 32-bit words incl. header pair */
	uint16_t arg1;
	uint8_t  module;
	uint8_t  arg3;
	uint16_t arg2;
} ipfw_insn;

struct ipfw_ioc_rule {
	uint16_t act_ofs;
	uint16_t cmd_len;
	uint16_t rulenum;
	uint8_t  set;
	uint8_t  insert;
	uint32_t sets;
	uint64_t pcnt;
	uint64_t bcnt;
	uint32_t timestamp;
	ipfw_insn cmd[1];
};

typedef struct {
	uint16_t opcode;
	uint16_t pad;
} ip_fw_x_header;

/* basic module opcodes (sys/net/ipfw3_basic/ip_fw3_basic.h) */
#define O_BASIC_ACCEPT	0
#define O_BASIC_DENY	1
#define O_BASIC_COUNT	2
#define O_BASIC_SKIPTO	3
#define O_BASIC_PROTO	10	/* filter: arg1 = protocol number */
#define LEN_OF_IPFWINSN 2

static int fw_sock;

static int
op_from_name(const char *n)
{
	if (!strcmp(n, "allow"))  return O_BASIC_ACCEPT;
	if (!strcmp(n, "deny"))   return O_BASIC_DENY;
	if (!strcmp(n, "count"))  return O_BASIC_COUNT;
	if (!strcmp(n, "skipto")) return O_BASIC_SKIPTO;
	return -1;
}

static const char *
op_name(uint8_t op)
{
	switch (op) {
	case O_BASIC_ACCEPT: return "allow";
	case O_BASIC_DENY:   return "deny";
	case O_BASIC_COUNT:  return "count";
	case O_BASIC_SKIPTO: return "skipto";
	case O_BASIC_PROTO:  return "proto";
	default:             return "?";
	}
}

static int
do_set_x(uint16_t opcode, void *payload, size_t paylen)
{
	unsigned char buf[2048];
	ip_fw_x_header *xh = (ip_fw_x_header *)buf;
	size_t len = sizeof(*xh) + paylen;

	memset(buf, 0, sizeof(buf));
	xh->opcode = opcode;
	memcpy(buf + sizeof(*xh), payload, paylen);
	return setsockopt(fw_sock, IPPROTO_IP, IP_FW_X, buf, len);
}

/* single-action rule: one 2-word instruction.
 * NOTE: struct ipfw_ioc_rule is padded to 48 bytes (align 8), so the
 * instruction array must be packed explicitly at offset 36. */
static int
do_add_action(uint16_t rulenum, uint8_t opcode, uint16_t arg1)
{
	unsigned char pay[64];
	struct ipfw_ioc_rule *ioc = (struct ipfw_ioc_rule *)pay;
	ipfw_insn *i0 = (ipfw_insn *)(pay + 36);

	memset(pay, 0, sizeof(pay));
	ioc->act_ofs = 0;
	ioc->cmd_len = LEN_OF_IPFWINSN;
	ioc->rulenum = rulenum;
	ioc->set = 0;
	ioc->insert = 0;
	i0->opcode = opcode;
	i0->len = LEN_OF_IPFWINSN;
	i0->module = 0;			/* MODULE_BASIC_ID */
	i0->arg1 = arg1;

	if (do_set_x(IP_FW_ADD, pay, 36 + 4 * LEN_OF_IPFWINSN) < 0) {
		fprintf(stderr, "add %u: setsockopt: %s\n", rulenum,
		    strerror(errno));
		return -1;
	}
	printf("added rule %u: %s %u\n", rulenum, op_name(opcode), arg1);
	return 0;
}

/* "proto <N> -> deny": filter insn + action insn = 4 words */
static int
do_add_protodeny(uint16_t rulenum, uint16_t proto)
{
	unsigned char pay[64];
	struct ipfw_ioc_rule *ioc = (struct ipfw_ioc_rule *)pay;
	ipfw_insn *i0 = (ipfw_insn *)(pay + 36);
	ipfw_insn *i1 = (ipfw_insn *)(pay + 44);

	memset(pay, 0, sizeof(pay));
	ioc->act_ofs = LEN_OF_IPFWINSN;		/* action = insn 1 */
	ioc->cmd_len = 2 * LEN_OF_IPFWINSN;
	ioc->rulenum = rulenum;
	ioc->set = 0;
	ioc->insert = 0;
	i0->opcode = O_BASIC_PROTO;
	i0->len = LEN_OF_IPFWINSN;
	i0->module = 0;
	i0->arg1 = proto;
	i1->opcode = O_BASIC_DENY;
	i1->len = LEN_OF_IPFWINSN;
	i1->module = 0;

	if (do_set_x(IP_FW_ADD, pay, 36 + 4 * ioc->cmd_len) < 0) {
		fprintf(stderr, "add %u: setsockopt: %s\n", rulenum,
		    strerror(errno));
		return -1;
	}
	printf("added rule %u: proto %u deny (4 words)\n", rulenum, proto);
	return 0;
}

static int
do_del(uint16_t rulenum)
{
	uint32_t rn = rulenum;

	if (do_set_x(IP_FW_DEL, &rn, sizeof(rn)) < 0) {
		fprintf(stderr, "del %u: setsockopt: %s\n", rulenum,
		    strerror(errno));
		return -1;
	}
	printf("deleted rule %u\n", rulenum);
	return 0;
}

static int
do_show(void)
{
	unsigned char buf[16384];
	socklen_t len = sizeof(buf);
	ip_fw_x_header *xh = (ip_fw_x_header *)buf;
	struct ipfw_ioc_rule *r;
	size_t off, total;

	memset(buf, 0, sizeof(buf));
	xh->opcode = IP_FW_GET;
	if (getsockopt(fw_sock, IPPROTO_IP, IP_FW_X, buf, &len) < 0) {
		fprintf(stderr, "show: getsockopt: %s\n", strerror(errno));
		return -1;
	}
	total = len;
	fprintf(stderr, "[show] returned len=%u\n", (unsigned)total);
	for (off = 0; off + 36 <= total;) {
		r = (struct ipfw_ioc_rule *)(buf + off);
		if (r->cmd_len < 2 || r->cmd_len > 250)
			break;
		printf("%6u pcnt=%-8llu bcnt=%-8llu cmd0=%s/%u arg1=%u "
		    "act_ofs=%u\n",
		    r->rulenum, (unsigned long long)r->pcnt,
		    (unsigned long long)r->bcnt,
		    op_name(r->cmd[0].opcode), r->cmd[0].len,
		    r->cmd[0].arg1, r->act_ofs);
		off += 36 + 4 * r->cmd_len;	/* IOC_RULESIZE */
	}
	return 0;
}

int
main(int argc, char **argv)
{
	uint16_t arg1 = 0;
	int op;

	if ((fw_sock = socket(AF_INET, SOCK_RAW, IPPROTO_RAW)) < 0) {
		fprintf(stderr, "socket: %s (need root)\n", strerror(errno));
		return 1;
	}

	if (argc < 2)
		goto usage;

	if (!strcmp(argv[1], "show"))
		return do_show() ? 1 : 0;
	if (!strcmp(argv[1], "flush")) {
		uint32_t zero = 0;
		if (do_set_x(IP_FW_FLUSH, &zero, sizeof(zero)) < 0) {
			fprintf(stderr, "flush: %s\n", strerror(errno));
			return 1;
		}
		printf("flushed\n");
		return 0;
	}
	if (!strcmp(argv[1], "del") && argc == 3)
		return do_del((uint16_t)strtoul(argv[2], NULL, 10)) ? 1 : 0;
	if (!strcmp(argv[1], "add") && (argc == 4 || argc == 5)) {
		if (!strcmp(argv[3], "protodeny") && argc == 5)
			return do_add_protodeny(
			    (uint16_t)strtoul(argv[2], NULL, 10),
			    (uint16_t)strtoul(argv[4], NULL, 10)) ? 1 : 0;
		op = op_from_name(argv[3]);
		if (op < 0)
			goto usage;
		if (argc == 5)
			arg1 = (uint16_t)strtoul(argv[4], NULL, 10);
		return do_add_action((uint16_t)strtoul(argv[2], NULL, 10), op,
		    arg1) ? 1 : 0;
	}
usage:
	fprintf(stderr,
	    "usage: %s add <rulenum> <allow|deny|count|skipto> [arg1]\n"
	    "       %s add <rulenum> protodeny <proto>\n"
	    "       %s del <rulenum>\n"
	    "       %s flush\n"
	    "       %s show\n",
	    argv[0], argv[0], argv[0], argv[0], argv[0]);
	return 1;
}