diff options
| -rw-r--r-- | data/templates/pmacct/override.conf.j2 | 4 | ||||
| -rw-r--r-- | data/templates/pmacct/uacctd.conf.j2 | 2 | ||||
| -rw-r--r-- | python/vyos/progressbar.py | 70 | ||||
| -rw-r--r-- | python/vyos/remote.py | 44 | ||||
| -rw-r--r-- | python/vyos/utils/io.py | 39 | ||||
| -rwxr-xr-x | smoketest/scripts/cli/test_interfaces_openvpn.py | 4 | ||||
| -rwxr-xr-x | src/conf_mode/flow_accounting_conf.py | 34 | ||||
| -rwxr-xr-x | src/system/uacctd_stop.py | 67 |
8 files changed, 200 insertions, 64 deletions
diff --git a/data/templates/pmacct/override.conf.j2 b/data/templates/pmacct/override.conf.j2 index 213569ddc..44a100bb6 100644 --- a/data/templates/pmacct/override.conf.j2 +++ b/data/templates/pmacct/override.conf.j2 @@ -9,9 +9,9 @@ ConditionPathExists=/run/pmacct/uacctd.conf EnvironmentFile= ExecStart= ExecStart={{ vrf_command }}/usr/sbin/uacctd -f /run/pmacct/uacctd.conf +ExecStop=/usr/libexec/vyos/system/uacctd_stop.py $MAINPID 60 WorkingDirectory= WorkingDirectory=/run/pmacct -PIDFile= -PIDFile=/run/pmacct/uacctd.pid Restart=always RestartSec=10 +KillMode=mixed diff --git a/data/templates/pmacct/uacctd.conf.j2 b/data/templates/pmacct/uacctd.conf.j2 index 1370f8121..aae0a0619 100644 --- a/data/templates/pmacct/uacctd.conf.j2 +++ b/data/templates/pmacct/uacctd.conf.j2 @@ -1,7 +1,7 @@ # Genereated from VyOS configuration daemonize: true promisc: false -pidfile: /run/pmacct/uacctd.pid +syslog: daemon uacctd_group: 2 uacctd_nl_size: 2097152 snaplen: {{ packet_length }} diff --git a/python/vyos/progressbar.py b/python/vyos/progressbar.py new file mode 100644 index 000000000..1793c445b --- /dev/null +++ b/python/vyos/progressbar.py @@ -0,0 +1,70 @@ +# Copyright 2023 VyOS maintainers and contributors <maintainers@vyos.io> +# +# This library is free software; you can redistribute it and/or +# modify it under the terms of the GNU Lesser General Public +# License as published by the Free Software Foundation; either +# version 2.1 of the License, or (at your option) any later version. +# +# This library is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +# Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public +# License along with this library. If not, see <http://www.gnu.org/licenses/>. + +import math +import os +import signal +import subprocess +import sys + +from vyos.utils.io import print_error + +class Progressbar: + def __init__(self, step=None): + self.total = 0.0 + self.step = step + def __enter__(self): + # Recalculate terminal width with every window resize. + signal.signal(signal.SIGWINCH, lambda signum, frame: self._update_cols()) + # Disable line wrapping to prevent the staircase effect. + subprocess.run(['tput', 'rmam'], check=False) + self._update_cols() + # Print an empty progressbar with entry. + self.progress(0, 1) + return self + def __exit__(self, exc_type, kexc_val, exc_tb): + # Revert to the default SIGWINCH handler (ie nothing). + signal.signal(signal.SIGWINCH, signal.SIG_DFL) + # Reenable line wrapping. + subprocess.run(['tput', 'smam'], check=False) + def _update_cols(self): + # `os.get_terminal_size()' is fast enough for our purposes. + self.col = max(os.get_terminal_size().columns - 15, 20) + def increment(self): + """ + Stateful progressbar taking the step fraction at init and no input at + callback (for FTP) + """ + if self.step: + if self.total < 1.0: + self.total += self.step + if self.total >= 1.0: + self.total = 1.0 + # Ignore superfluous calls caused by fuzzy FTP size calculations. + self.step = None + self.progress(self.total, 1.0) + def progress(self, done, total): + """ + Stateless progressbar taking no input at init and current progress with + final size at callback (for SSH) + """ + if done <= total: + length = math.ceil(self.col * done / total) + percentage = str(math.ceil(100 * done / total)).rjust(3) + # Carriage return at the end will make sure the line will get overwritten. + print_error(f'[{length * "#"}{(self.col - length) * "_"}] {percentage}%', end='\r') + # Print a newline to make sure the full progressbar doesn't get overwritten by the next line. + if done == total: + print_error() diff --git a/python/vyos/remote.py b/python/vyos/remote.py index cf731c881..1ca8a9530 100644 --- a/python/vyos/remote.py +++ b/python/vyos/remote.py @@ -32,9 +32,8 @@ from requests import Session from requests.adapters import HTTPAdapter from requests.packages.urllib3 import PoolManager +from vyos.progressbar import Progressbar from vyos.utils.io import ask_yes_no -from vyos.utils.io import make_incremental_progressbar -from vyos.utils.io import make_progressbar from vyos.utils.io import print_error from vyos.utils.misc import begin from vyos.utils.process import cmd @@ -131,16 +130,16 @@ class FtpC: if self.secure: conn.prot_p() # Almost all FTP servers support the `SIZE' command. + size = conn.size(self.path) if self.check_space: - check_storage(path, conn.size(self.path)) + check_storage(path, size) # No progressbar if we can't determine the size or if the file is too small. if self.progressbar and size and size > CHUNK_SIZE: - progress = make_incremental_progressbar(CHUNK_SIZE / size) - next(progress) - callback = lambda block: begin(f.write(block), next(progress)) + with Progressbar(CHUNK_SIZE / size) as p: + callback = lambda block: begin(f.write(block), p.increment()) + conn.retrbinary('RETR ' + self.path, callback, CHUNK_SIZE) else: - callback = f.write - conn.retrbinary('RETR ' + self.path, callback, CHUNK_SIZE) + conn.retrbinary('RETR ' + self.path, f.write, CHUNK_SIZE) def upload(self, location: str): size = os.path.getsize(location) @@ -150,12 +149,10 @@ class FtpC: if self.secure: conn.prot_p() if self.progressbar and size and size > CHUNK_SIZE: - progress = make_incremental_progressbar(CHUNK_SIZE / size) - next(progress) - callback = lambda block: next(progress) + with Progressbar(CHUNK_SIZE / size) as p: + conn.storbinary('STOR ' + self.path, f, CHUNK_SIZE, lambda block: p.increment()) else: - callback = None - conn.storbinary('STOR ' + self.path, f, CHUNK_SIZE, callback) + conn.storbinary('STOR ' + self.path, f, CHUNK_SIZE) class SshC: known_hosts = os.path.expanduser('~/.ssh/known_hosts') @@ -190,14 +187,16 @@ class SshC: return ssh def download(self, location: str): - callback = make_progressbar() if self.progressbar else None with self._establish() as ssh, ssh.open_sftp() as sftp: if self.check_space: check_storage(location, sftp.stat(self.path).st_size) - sftp.get(self.path, location, callback=callback) + if self.progressbar: + with Progressbar() as p: + sftp.get(self.path, location, callback=p.progress) + else: + sftp.get(self.path, location) def upload(self, location: str): - callback = make_progressbar() if self.progressbar else None with self._establish() as ssh, ssh.open_sftp() as sftp: try: # If the remote path is a directory, use the original filename. @@ -210,7 +209,11 @@ class SshC: except IOError: path = self.path finally: - sftp.put(location, path, callback=callback) + if self.progressbar: + with Progressbar() as p: + sftp.put(location, path, callback=p.progress) + else: + sftp.put(location, path) class HttpC: @@ -264,10 +267,9 @@ class HttpC: with s.get(final_urlstring, stream=True, timeout=self.timeout) as r, open(location, 'wb') as f: if self.progressbar and size: - progress = make_incremental_progressbar(CHUNK_SIZE / size) - next(progress) - for chunk in iter(lambda: begin(next(progress), r.raw.read(CHUNK_SIZE)), b''): - f.write(chunk) + with Progressbar(CHUNK_SIZE / size) as p: + for chunk in iter(lambda: begin(p.increment(), r.raw.read(CHUNK_SIZE)), b''): + f.write(chunk) else: # We'll try to stream the download directly with `copyfileobj()` so that large # files (like entire VyOS images) don't occupy much memory. diff --git a/python/vyos/utils/io.py b/python/vyos/utils/io.py index 843494855..5fffa62f8 100644 --- a/python/vyos/utils/io.py +++ b/python/vyos/utils/io.py @@ -24,45 +24,6 @@ def print_error(str='', end='\n'): sys.stderr.write(end) sys.stderr.flush() -def make_progressbar(): - """ - Make a procedure that takes two arguments `done` and `total` and prints a - progressbar based on the ratio thereof, whose length is determined by the - width of the terminal. - """ - import shutil, math - col, _ = shutil.get_terminal_size() - col = max(col - 15, 20) - def print_progressbar(done, total): - if done <= total: - increment = total / col - length = math.ceil(done / increment) - percentage = str(math.ceil(100 * done / total)).rjust(3) - print_error(f'[{length * "#"}{(col - length) * "_"}] {percentage}%', '\r') - # Print a newline so that the subsequent prints don't overwrite the full bar. - if done == total: - print_error() - return print_progressbar - -def make_incremental_progressbar(increment: float): - """ - Make a generator that displays a progressbar that grows monotonically with - every iteration. - First call displays it at 0% and every subsequent iteration displays it - at `increment` increments where 0.0 < `increment` < 1.0. - Intended for FTP and HTTP transfers with stateless callbacks. - """ - print_progressbar = make_progressbar() - total = 0.0 - while total < 1.0: - print_progressbar(total, 1.0) - yield - total += increment - print_progressbar(1, 1) - # Ignore further calls. - while True: - yield - def ask_input(question, default='', numeric_only=False, valid_responses=[]): question_out = question if default: diff --git a/smoketest/scripts/cli/test_interfaces_openvpn.py b/smoketest/scripts/cli/test_interfaces_openvpn.py index 4a7e2418c..66c348976 100755 --- a/smoketest/scripts/cli/test_interfaces_openvpn.py +++ b/smoketest/scripts/cli/test_interfaces_openvpn.py @@ -506,11 +506,13 @@ class TestInterfacesOpenVPN(VyOSUnitTestSHIM.TestCase): interface = 'vtun5001' path = base_path + [interface] + encryption_cipher = 'aes256' self.cli_set(path + ['mode', 'site-to-site']) self.cli_set(path + ['local-address', '10.0.0.2']) self.cli_set(path + ['remote-address', '192.168.0.3']) self.cli_set(path + ['shared-secret-key', 'ovpn_test']) + self.cli_set(path + ['encryption', 'cipher', encryption_cipher]) self.cli_commit() @@ -548,6 +550,7 @@ class TestInterfacesOpenVPN(VyOSUnitTestSHIM.TestCase): port = '' local_address = '' remote_address = '' + encryption_cipher = 'aes256' for ii in num_range: interface = f'vtun{ii}' @@ -571,6 +574,7 @@ class TestInterfacesOpenVPN(VyOSUnitTestSHIM.TestCase): self.cli_set(path + ['remote-port', port]) self.cli_set(path + ['shared-secret-key', 'ovpn_test']) self.cli_set(path + ['remote-address', remote_address]) + self.cli_set(path + ['encryption', 'cipher', encryption_cipher]) self.cli_set(path + ['vrf', vrf_name]) self.cli_commit() diff --git a/src/conf_mode/flow_accounting_conf.py b/src/conf_mode/flow_accounting_conf.py index 81ee39df1..206f513c8 100755 --- a/src/conf_mode/flow_accounting_conf.py +++ b/src/conf_mode/flow_accounting_conf.py @@ -28,6 +28,7 @@ from vyos.ifconfig import Section from vyos.template import render from vyos.utils.process import call from vyos.utils.process import cmd +from vyos.utils.process import run from vyos.utils.network import is_addr_assigned from vyos import ConfigError from vyos import airbag @@ -116,6 +117,30 @@ def _nftables_config(configured_ifaces, direction, length=None): cmd(command, raising=ConfigError) +def _nftables_trigger_setup(operation: str) -> None: + """Add a dummy rule to unlock the main pmacct loop with a packet-trigger + + Args: + operation (str): 'add' or 'delete' a trigger + """ + # check if a chain exists + table_exists = False + if run('nft -snj list table ip pmacct') == 0: + table_exists = True + + if operation == 'delete' and table_exists: + nft_cmd: str = 'nft delete table ip pmacct' + cmd(nft_cmd, raising=ConfigError) + if operation == 'add' and not table_exists: + nft_cmds: list[str] = [ + 'nft add table ip pmacct', + 'nft add chain ip pmacct pmacct_out { type filter hook output priority raw - 50 \\; policy accept \\; }', + 'nft add rule ip pmacct pmacct_out oif lo ip daddr 127.0.254.0 counter log group 2 snaplen 1 queue-threshold 0 comment NFLOG_TRIGGER' + ] + for nft_cmd in nft_cmds: + cmd(nft_cmd, raising=ConfigError) + + def get_config(config=None): if config: conf = config @@ -252,7 +277,6 @@ def generate(flow_config): call('systemctl daemon-reload') def apply(flow_config): - action = 'restart' # Check if flow-accounting was removed and define command if not flow_config: _nftables_config([], 'ingress') @@ -262,6 +286,10 @@ def apply(flow_config): call(f'systemctl stop {systemd_service}') if os.path.exists(uacctd_conf_path): os.unlink(uacctd_conf_path) + + # must be done after systemctl + _nftables_trigger_setup('delete') + return # Start/reload flow-accounting daemon @@ -277,6 +305,10 @@ def apply(flow_config): else: _nftables_config([], 'egress') + # add a trigger for signal processing + _nftables_trigger_setup('add') + + if __name__ == '__main__': try: config = get_config() diff --git a/src/system/uacctd_stop.py b/src/system/uacctd_stop.py new file mode 100755 index 000000000..7fbac0566 --- /dev/null +++ b/src/system/uacctd_stop.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +# +# Copyright (C) 2023 VyOS maintainers and contributors +# +# This program is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License version 2 or later as +# published by the Free Software Foundation. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see <http://www.gnu.org/licenses/>. + +# Control pmacct daemons in a tricky way. +# Pmacct has signal processing in a main loop, together with packet +# processing. Because of this, while it is waiting for packets, it cannot +# handle the control signal. We need to start the systemctl command and then +# send some packets to pmacct to wake it up + +from argparse import ArgumentParser +from socket import socket +from sys import exit +from time import sleep + +from psutil import Process + + +def stop_process(pid: int, timeout: int) -> None: + """Send a signal to uacctd + and then send packets to special address predefined in a firewall + to unlock main loop in uacctd and finish the process properly + + Args: + pid (int): uacctd PID + timeout (int): seconds to wait for a process end + """ + # find a process + uacctd = Process(pid) + uacctd.terminate() + + # create a socket + trigger = socket() + + first_cycle: bool = True + while uacctd.is_running() and timeout: + trigger.sendto(b'WAKEUP', ('127.0.254.0', 0)) + # do not sleep during first attempt + if not first_cycle: + sleep(1) + timeout -= 1 + first_cycle = False + + +if __name__ == '__main__': + parser = ArgumentParser() + parser.add_argument('process_id', + type=int, + help='PID file of uacctd core process') + parser.add_argument('timeout', + type=int, + help='time to wait for process end') + args = parser.parse_args() + stop_process(args.process_id, args.timeout) + exit() |