mirror of
https://gitlab.nic.cz/knot/knot-dns.git
synced 2026-02-03 18:49:28 -05:00
Convert script rom Bash to Python3
This commit is contained in:
Jan Hák
parent
94f8e4ace3
commit
2c2e2762b1
3 changed files with 122 additions and 216 deletions
90
kxdpconf.py
Executable file
90
kxdpconf.py
Executable file
|
|
@ -0,0 +1,90 @@
|
|||
#!/usr/bin/env python3
|
||||
|
||||
# Dependencies:
|
||||
# pip install argparse enum
|
||||
# sudo apt install python3-bpfcc
|
||||
|
||||
from argparse import ArgumentParser, ArgumentTypeError
|
||||
from bcc import BPF
|
||||
from enum import IntFlag
|
||||
|
||||
class Flags(IntFlag):
|
||||
ENABLED = 1 << 0
|
||||
UDP = 1 << 1
|
||||
TCP = 1 << 2
|
||||
QUIC = 1 << 3
|
||||
|
||||
bpf_template = r'''
|
||||
struct knot_xdp_opts {
|
||||
u16 flags;
|
||||
u16 udp_port;
|
||||
u16 quic_port;
|
||||
} __attribute__((packed));
|
||||
BPF_TABLE_PINNED("array", u32, struct knot_xdp_opts, opts_map, 256, "/sys/fs/bpf/knot/opts_%s");
|
||||
'''
|
||||
# Parse arguments
|
||||
def parse_u16(value):
|
||||
try:
|
||||
unsigned_short_value = int(value)
|
||||
if 0 <= unsigned_short_value <= 65535:
|
||||
return unsigned_short_value
|
||||
else:
|
||||
raise ArgumentTypeError("Value must be an unsigned short (0 to 65535)")
|
||||
except ValueError:
|
||||
raise ArgumentTypeError("Invalid unsigned short value")
|
||||
|
||||
def parse_on_off(value):
|
||||
if value.lower() == 'on':
|
||||
return True
|
||||
elif value.lower() == 'off':
|
||||
return False
|
||||
else:
|
||||
raise ArgumentTypeError("Invalid value for --udp. Use 'on', 'off', or leave it unspecified.")
|
||||
|
||||
parser = ArgumentParser(description="Your script description.")
|
||||
parser.add_argument("interfaces", nargs="+", help="List of interfaces")
|
||||
parser.add_argument("-u", "--udp", type=parse_on_off, default=None, choices=["on", "off"], help="Enable/disable UDP")
|
||||
parser.add_argument("-t", "--tcp", type=parse_on_off, default=None, choices=["on", "off"], help="Enable/disable TCP")
|
||||
parser.add_argument("-q", "--quic", type=parse_on_off, default=None, choices=["on", "off"], help="Enable/disable QUIC")
|
||||
parser.add_argument("-p", "--port", type=parse_u16, default=None, help="Specify UDP/TCP port")
|
||||
parser.add_argument("-Q", "--quic_port", type=parse_u16, default=None, help="Specify QUIC port")
|
||||
parser.add_argument("-V", "--version", action="version", version="Your Script Version 0.0.1")
|
||||
args = parser.parse_args()
|
||||
|
||||
# For each interface
|
||||
for iface in args.interfaces:
|
||||
# Compile BPF filter
|
||||
bpf = bpf_template % iface
|
||||
try:
|
||||
xdp_stats_map = BPF(text=bpf, cflags=['-w']).get_table('opts_map')
|
||||
except Exception as e:
|
||||
print(f"Error: Unable to locate 'opts_map' related to '%s'. Is server running?" % iface)
|
||||
continue
|
||||
|
||||
# For each key-value pair in 'opts_map'
|
||||
for k, v in xdp_stats_map.items():
|
||||
if (not v.flags & Flags.ENABLED):
|
||||
continue
|
||||
|
||||
if args.udp is True:
|
||||
v.flags |= Flags.UDP
|
||||
elif args.udp is False:
|
||||
v.flags ^= Flags.UDP
|
||||
|
||||
if args.tcp is True:
|
||||
v.flags |= Flags.TCP
|
||||
elif args.tcp is False:
|
||||
v.flags ^= Flags.TCP
|
||||
|
||||
if args.quic is True:
|
||||
v.flags |= Flags.QUIC
|
||||
elif args.quic is False:
|
||||
v.flags ^= Flags.QUIC
|
||||
|
||||
if args.port is not None:
|
||||
v.udp_port = args.port
|
||||
|
||||
if args.quic_port is not None:
|
||||
v.quic_port = args.quic_port
|
||||
|
||||
xdp_stats_map[k] = v
|
||||
210
kxdpconf.sh
210
kxdpconf.sh
|
|
@ -1,210 +0,0 @@
|
|||
#!/usr/bin/env bash
|
||||
#
|
||||
# This is a rather minimal example Argbash potential
|
||||
# Example taken from http://argbash.readthedocs.io/en/stable/example.html
|
||||
#
|
||||
# ARG_POSITIONAL_INF([interface],[Enumeration of network interfaces],[1])
|
||||
# ARG_OPTIONAL_SINGLE([flags],[f],[Flags])
|
||||
# ARG_OPTIONAL_SINGLE([port],[p],[UDP/TCP port])
|
||||
# ARG_OPTIONAL_SINGLE([quic_port],[q],[QUIC port])
|
||||
# ARG_HELP([The general script's help msg])
|
||||
# ARG_VERSION([0.0.1])
|
||||
# ARGBASH_GO()
|
||||
# needed because of Argbash --> m4_ignore([
|
||||
### START OF CODE GENERATED BY Argbash v2.9.0 one line above ###
|
||||
# Argbash is a bash code generator used to get arguments parsing right.
|
||||
# Argbash is FREE SOFTWARE, see https://argbash.io for more info
|
||||
# Generated online by https://argbash.io/generate
|
||||
|
||||
|
||||
die()
|
||||
{
|
||||
local _ret="${2:-1}"
|
||||
test "${_PRINT_HELP:-no}" = yes && print_help >&2
|
||||
echo "$1" >&2
|
||||
exit "${_ret}"
|
||||
}
|
||||
|
||||
|
||||
begins_with_short_option()
|
||||
{
|
||||
local first_option all_short_options='fpqhv'
|
||||
first_option="${1:0:1}"
|
||||
test "$all_short_options" = "${all_short_options/$first_option/}" && return 1 || return 0
|
||||
}
|
||||
|
||||
# THE DEFAULTS INITIALIZATION - POSITIONALS
|
||||
_positionals=()
|
||||
_arg_interface=('' )
|
||||
# THE DEFAULTS INITIALIZATION - OPTIONALS
|
||||
_arg_flags=
|
||||
_arg_port=
|
||||
_arg_quic_port=
|
||||
|
||||
|
||||
print_help()
|
||||
{
|
||||
printf '%s\n' "The general script's help msg"
|
||||
printf 'Usage: %s [-f|--flags <arg>] [-p|--port <arg>] [-q|--quic_port <arg>] [-h|--help] [-v|--version] <interface-1> [<interface-2>] ... [<interface-n>] ...\n' "$0"
|
||||
printf '\t%s\n' "<interface>: Enumeration of network interfaces"
|
||||
printf '\t%s\n' "-f, --flags: Flags (no default)"
|
||||
printf '\t%s\n' "-p, --port: UDP/TCP port (no default)"
|
||||
printf '\t%s\n' "-q, --quic_port: QUIC port (no default)"
|
||||
printf '\t%s\n' "-h, --help: Prints help"
|
||||
printf '\t%s\n' "-v, --version: Prints version"
|
||||
}
|
||||
|
||||
|
||||
parse_commandline()
|
||||
{
|
||||
_positionals_count=0
|
||||
while test $# -gt 0
|
||||
do
|
||||
_key="$1"
|
||||
case "$_key" in
|
||||
-f|--flags)
|
||||
test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
|
||||
_arg_flags="$2"
|
||||
shift
|
||||
;;
|
||||
--flags=*)
|
||||
_arg_flags="${_key##--flags=}"
|
||||
;;
|
||||
-f*)
|
||||
_arg_flags="${_key##-f}"
|
||||
;;
|
||||
-p|--port)
|
||||
test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
|
||||
_arg_port="$2"
|
||||
shift
|
||||
;;
|
||||
--port=*)
|
||||
_arg_port="${_key##--port=}"
|
||||
;;
|
||||
-p*)
|
||||
_arg_port="${_key##-p}"
|
||||
;;
|
||||
-q|--quic_port)
|
||||
test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
|
||||
_arg_quic_port="$2"
|
||||
shift
|
||||
;;
|
||||
--quic_port=*)
|
||||
_arg_quic_port="${_key##--quic_port=}"
|
||||
;;
|
||||
-q*)
|
||||
_arg_quic_port="${_key##-q}"
|
||||
;;
|
||||
-h|--help)
|
||||
print_help
|
||||
exit 0
|
||||
;;
|
||||
-h*)
|
||||
print_help
|
||||
exit 0
|
||||
;;
|
||||
-v|--version)
|
||||
0.0.1
|
||||
exit 0
|
||||
;;
|
||||
-v*)
|
||||
0.0.1
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
_last_positional="$1"
|
||||
_positionals+=("$_last_positional")
|
||||
_positionals_count=$((_positionals_count + 1))
|
||||
;;
|
||||
esac
|
||||
shift
|
||||
done
|
||||
}
|
||||
|
||||
|
||||
handle_passed_args_count()
|
||||
{
|
||||
local _required_args_string="'interface'"
|
||||
test "${_positionals_count}" -ge 1 || _PRINT_HELP=yes die "FATAL ERROR: Not enough positional arguments - we require at least 1 (namely: $_required_args_string), but got only ${_positionals_count}." 1
|
||||
}
|
||||
|
||||
|
||||
assign_positional_args()
|
||||
{
|
||||
local _positional_name _shift_for=$1
|
||||
_positional_names="_arg_interface "
|
||||
_our_args=$((${#_positionals[@]} - 1))
|
||||
for ((ii = 0; ii < _our_args; ii++))
|
||||
do
|
||||
_positional_names="$_positional_names _arg_interface[$((ii + 1))]"
|
||||
done
|
||||
|
||||
shift "$_shift_for"
|
||||
for _positional_name in ${_positional_names}
|
||||
do
|
||||
test $# -gt 0 || break
|
||||
eval "$_positional_name=\${1}" || die "Error during argument parsing, possibly an Argbash bug." 1
|
||||
shift
|
||||
done
|
||||
}
|
||||
|
||||
parse_commandline "$@"
|
||||
handle_passed_args_count
|
||||
assign_positional_args 1 "${_positionals[@]}"
|
||||
|
||||
# OTHER STUFF GENERATED BY Argbash
|
||||
|
||||
### END OF CODE GENERATED BY Argbash (sortof) ### ])
|
||||
# [ <-- needed because of Argbash
|
||||
|
||||
dump_map()
|
||||
{
|
||||
bpftool map dump pinned $1 | head -n -1
|
||||
}
|
||||
|
||||
parse_key_val()
|
||||
{
|
||||
awk '{ if ($7 == "00" && $8 == "00" && $9 == "00" && $10 == "00" && $11 == "00" && $12 == "00") next; print $2, $3, $4, $5 "\t" $7, $8, $9, $10, $11, $12}'
|
||||
}
|
||||
|
||||
to_hex()
|
||||
{
|
||||
printf "%04x" $1 | sed -E 's/(..)(..)/\2 \1/'
|
||||
}
|
||||
|
||||
set_port()
|
||||
{
|
||||
[ $2 -gt 65535 ] && die "ERROR: Port has to be in range 0 - 65535" 1
|
||||
read -ra output_array <<< "$1"
|
||||
read -ra new_val_array <<< $( to_hex "$2" )
|
||||
|
||||
printf "%s %s %s %s %s %s" ${output_array[0]} ${output_array[1]} ${new_val_array[0]} ${new_val_array[1]} ${output_array[4]} ${output_array[5]}
|
||||
|
||||
}
|
||||
|
||||
set_quic_port()
|
||||
{
|
||||
[ $2 -gt 65535 ] && die "ERROR: Port has to be in range 0 - 65535" 1
|
||||
read -ra output_array <<< "$1"
|
||||
read -ra new_val_array <<< $( to_hex "$2" )
|
||||
|
||||
printf "%s %s %s %s %s %s" ${output_array[0]} ${output_array[1]} ${output_array[2]} ${output_array[3]} ${new_val_array[0]} ${new_val_array[1]}
|
||||
|
||||
}
|
||||
|
||||
for iface in "${_arg_interface[@]}"; do
|
||||
[ -z "$iface" ] && continue
|
||||
MAP_PATH=$(printf "/sys/fs/bpf/knot/%s/opts_map" $iface)
|
||||
dump_map $MAP_PATH | parse_key_val | while IFS=$'\t' read -r key value; do
|
||||
new_value="$value"
|
||||
|
||||
# [ -n "$_arg_flags" ] && echo "changing flags"
|
||||
[ -n "$_arg_port" ] && new_value=$( set_port "$new_value" "$_arg_port" )
|
||||
[ -n "$_arg_quic_port" ] && new_value=$( set_quic_port "$new_value" "$_arg_quic_port" )
|
||||
|
||||
[ "$value" == "$new_value" ] && continue
|
||||
bpftool map update pinned $MAP_PATH key hex $key value hex $new_value
|
||||
done
|
||||
done
|
||||
|
||||
# ] <-- needed because of Argbash
|
||||
|
|
@ -35,8 +35,6 @@ static inline bool IS_ERR_OR_NULL(const void *ptr)
|
|||
return (ptr == NULL) || (unsigned long)ptr >= (unsigned long)-4095;
|
||||
}
|
||||
|
||||
#include <linux/limits.h>
|
||||
|
||||
static int prog_load(const char *iface, struct bpf_object **pobj, int *prog_fd)
|
||||
{
|
||||
struct bpf_program *prog, *first_prog = NULL;
|
||||
|
|
@ -65,11 +63,32 @@ static int prog_load(const char *iface, struct bpf_object **pobj, int *prog_fd)
|
|||
return KNOT_EINVAL;
|
||||
}
|
||||
|
||||
char pin_dir[PATH_MAX];
|
||||
snprintf(pin_dir, sizeof(pin_dir), "/sys/fs/bpf/knot/%s", iface);
|
||||
ret = bpf_object__pin_maps(obj, pin_dir);
|
||||
struct bpf_map *opts_map = bpf_object__find_map_by_name(obj, "opts_map");
|
||||
if (opts_map == NULL) {
|
||||
bpf_object__close(obj);
|
||||
return KNOT_EINVAL;
|
||||
}
|
||||
|
||||
char pin_dir[64];
|
||||
snprintf(pin_dir, sizeof(pin_dir), "/sys/fs/bpf/knot/opts_%s", iface);
|
||||
ret = bpf_map__set_pin_path(opts_map, pin_dir);
|
||||
if (ret != 0) {
|
||||
|
||||
bpf_object__close(obj);
|
||||
return KNOT_EINVAL;
|
||||
}
|
||||
|
||||
/* Existing/previous XDP prog might not have cleaned up */
|
||||
if (access(pin_dir, F_OK) != -1 ) {
|
||||
ret = unlink(pin_dir);
|
||||
if (ret) {
|
||||
bpf_object__close(obj);
|
||||
return knot_map_errno();
|
||||
}
|
||||
}
|
||||
ret = bpf_map__pin(opts_map, NULL);
|
||||
if (ret != 0) {
|
||||
bpf_object__close(obj);
|
||||
return ret;
|
||||
}
|
||||
|
||||
*pobj = obj;
|
||||
|
|
@ -128,6 +147,13 @@ static int ensure_prog(struct kxsk_iface *iface, bool overwrite, bool generic_xd
|
|||
|
||||
static void unget_bpf_maps(struct kxsk_iface *iface)
|
||||
{
|
||||
/*
|
||||
* TODO This line does nothing.
|
||||
* `knotd` currently lacks the ability to unpin maps due to
|
||||
* the dropping of user capabilities post-initialization, resulting in
|
||||
* an EACCES error upon attempted call.
|
||||
*/
|
||||
bpf_object__unpin_maps(iface->prog_obj, NULL);
|
||||
if (iface->opts_map_fd >= 0) {
|
||||
close(iface->opts_map_fd);
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue