You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
106 lines
3.0 KiB
Python
106 lines
3.0 KiB
Python
import argparse
|
|
import signal
|
|
import time
|
|
import sys
|
|
from pydantic import BaseModel, StrictStr, StrictInt, root_validator
|
|
from typing_extensions import Literal
|
|
from .mtu_finder import MTUFinder
|
|
|
|
|
|
def signal_handler(sig, frame):
|
|
"""Handle ctrl+c interrupt.
|
|
|
|
Without this handler, everytime a ctrl+c interrupt is received, the server shutdowns and
|
|
proceeds to the next iteration in the loop rather than exiting the program altogether.
|
|
"""
|
|
print("************Received CTRL-C. Will exit in 1 second************")
|
|
time.sleep(1)
|
|
sys.exit(0)
|
|
|
|
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
|
|
|
|
class ArgsModel(BaseModel):
|
|
mode: Literal["server", "peer"]
|
|
mtu_min: int
|
|
mtu_max: int
|
|
mtu_step: int
|
|
|
|
server_ip: StrictStr
|
|
server_port: int = 5000
|
|
|
|
interface: StrictStr = "wg0"
|
|
conf_file: StrictStr = "/etc/wireguard/wg0.conf"
|
|
|
|
@root_validator(pre=False)
|
|
def validate(cls, values):
|
|
"""Generic validations."""
|
|
mtu_min, mtu_max, mtu_step = (
|
|
values.get("mtu_min", None),
|
|
values.get("mtu_max", None),
|
|
values.get("mtu_step", None),
|
|
)
|
|
|
|
if not (1280 <= mtu_min <= 1500):
|
|
raise ValueError(f"mtu_min: {mtu_min} must be in range [1280, 1500].")
|
|
|
|
if not (1280 <= mtu_max <= 1500):
|
|
raise ValueError(f"mtu_max: {mtu_max} must be in range [1280, 1500].")
|
|
|
|
if not (mtu_min <= mtu_max):
|
|
raise ValueError(f"mtu_min: {mtu_min} must be less than or equal to mtu_max: {mtu_max}")
|
|
|
|
return values
|
|
|
|
class Config:
|
|
orm_mode = True
|
|
|
|
|
|
def setup_args():
|
|
"""Setup args."""
|
|
parser = argparse.ArgumentParser(
|
|
description="nr-wg-mtu-finder - Helps find the optimal Wireguard MTU between Server and Peer."
|
|
)
|
|
parser.add_argument(
|
|
"--mode",
|
|
help=(
|
|
"Mode is 'server' if you are running this script on the WG Server, "
|
|
"else the mode is 'peer' if you are running this script on the WG Peer."
|
|
),
|
|
required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--mtu-min", help="Min MTU.", required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--mtu-max", help="Max MTU.", required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--mtu-step", help="By how much to increment the MTU between loops.", required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--server-ip", help="The IP address of the WG server and flask server.", required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--server-port", help="The port for the flask server.", required=False, default=5000,
|
|
)
|
|
parser.add_argument(
|
|
"--interface", help="The WG interface name. Default: 'wg0'", required=False, default="wg0"
|
|
)
|
|
parser.add_argument(
|
|
"--conf-file",
|
|
help="The WG interface name. Default: '/etc/wireguard/wg0.conf'",
|
|
required=False,
|
|
default="/etc/wireguard/wg0.conf",
|
|
)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def run():
|
|
args = setup_args()
|
|
args = ArgsModel.from_orm(args)
|
|
|
|
MTUFinder(**args.dict())
|