diff options
Diffstat (limited to 'python/vyos/util.py')
| -rw-r--r-- | python/vyos/util.py | 174 |
1 files changed, 145 insertions, 29 deletions
diff --git a/python/vyos/util.py b/python/vyos/util.py index 2a3f6a228..59f9f1c44 100644 --- a/python/vyos/util.py +++ b/python/vyos/util.py @@ -1,4 +1,4 @@ -# Copyright 2020 VyOS maintainers and contributors <maintainers@vyos.io> +# Copyright 2020-2021 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 @@ -22,25 +22,13 @@ import sys # where it is used so it is as local as possible to the execution # - -def _need_sudo(command): - return os.path.basename(command.split()[0]) in ('systemctl', ) - - -def _add_sudo(command): - if _need_sudo(command): - return 'sudo ' + command - return command - - from subprocess import Popen from subprocess import PIPE from subprocess import STDOUT from subprocess import DEVNULL - def popen(command, flag='', shell=None, input=None, timeout=None, env=None, - stdout=PIPE, stderr=PIPE, decode='utf-8', autosudo=True): + stdout=PIPE, stderr=PIPE, decode='utf-8'): """ popen is a wrapper helper aound subprocess.Popen with it default setting it will return a tuple (out, err) @@ -79,9 +67,6 @@ def popen(command, flag='', shell=None, input=None, timeout=None, env=None, if not debug.enabled(flag): flag = 'command' - if autosudo: - command = _add_sudo(command) - cmd_msg = f"cmd '{command}'" debug.message(cmd_msg, flag) @@ -98,11 +83,8 @@ def popen(command, flag='', shell=None, input=None, timeout=None, env=None, stdin = PIPE input = input.encode() if type(input) is str else input - p = Popen( - command, - stdin=stdin, stdout=stdout, stderr=stderr, - env=env, shell=use_shell, - ) + p = Popen(command, stdin=stdin, stdout=stdout, stderr=stderr, + env=env, shell=use_shell) pipe = p.communicate(input, timeout) @@ -135,7 +117,7 @@ def popen(command, flag='', shell=None, input=None, timeout=None, env=None, def run(command, flag='', shell=None, input=None, timeout=None, env=None, - stdout=DEVNULL, stderr=PIPE, decode='utf-8', autosudo=True): + stdout=DEVNULL, stderr=PIPE, decode='utf-8'): """ A wrapper around popen, which discard the stdout and will return the error code of a command @@ -151,8 +133,8 @@ def run(command, flag='', shell=None, input=None, timeout=None, env=None, def cmd(command, flag='', shell=None, input=None, timeout=None, env=None, - stdout=PIPE, stderr=PIPE, decode='utf-8', autosudo=True, - raising=None, message='', expect=[0]): + stdout=PIPE, stderr=PIPE, decode='utf-8', raising=None, message='', + expect=[0]): """ A wrapper around popen, which returns the stdout and will raise the error code of a command @@ -183,7 +165,7 @@ def cmd(command, flag='', shell=None, input=None, timeout=None, env=None, def call(command, flag='', shell=None, input=None, timeout=None, env=None, - stdout=PIPE, stderr=PIPE, decode='utf-8', autosudo=True): + stdout=PIPE, stderr=PIPE, decode='utf-8'): """ A wrapper around popen, which print the stdout and will return the error code of a command @@ -239,7 +221,6 @@ def write_file(fname, data, defaultonfailure=None, user=None, group=None): return defaultonfailure raise e - def read_json(fname, defaultonfailure=None): """ read and json decode the content of a file @@ -459,7 +440,6 @@ def process_running(pid_file): pid = f.read().strip() return pid_exists(int(pid)) - def process_named_running(name): """ Checks if process with given name is running and returns its PID. If Process is not running, return None @@ -470,7 +450,6 @@ def process_named_running(name): return p.pid return None - def seconds_to_human(s, separator=""): """ Converts number of seconds passed to a human-readable interval such as 1w4d18h35m59s @@ -525,6 +504,46 @@ def file_is_persistent(path): absolute = os.path.abspath(os.path.dirname(path)) return re.match(location,absolute) +def wait_for_inotify(file_path, pre_hook=None, event_type=None, timeout=None, sleep_interval=0.1): + """ Waits for an inotify event to occur """ + if not os.path.dirname(file_path): + raise ValueError( + "File path {} does not have a directory part (required for inotify watching)".format(file_path)) + if not os.path.basename(file_path): + raise ValueError( + "File path {} does not have a file part, do not know what to watch for".format(file_path)) + + from inotify.adapters import Inotify + from time import time + from time import sleep + + time_start = time() + + i = Inotify() + i.add_watch(os.path.dirname(file_path)) + + if pre_hook: + pre_hook() + + for event in i.event_gen(yield_nones=True): + if (timeout is not None) and ((time() - time_start) > timeout): + # If the function didn't return until this point, + # the file failed to have been written to and closed within the timeout + raise OSError("Waiting for file {} to be written has failed".format(file_path)) + + # Most such events don't take much time, so it's better to check right away + # and sleep later. + if event is not None: + (_, type_names, path, filename) = event + if filename == os.path.basename(file_path): + if event_type in type_names: + return + sleep(sleep_interval) + +def wait_for_file_write_complete(file_path, pre_hook=None, timeout=None, sleep_interval=0.1): + """ Waits for a process to close a file after opening it in write mode. """ + wait_for_inotify(file_path, + event_type='IN_CLOSE_WRITE', pre_hook=pre_hook, timeout=timeout, sleep_interval=sleep_interval) def commit_in_progress(): """ Not to be used in normal op mode scripts! """ @@ -571,6 +590,25 @@ def wait_for_commit_lock(): while commit_in_progress(): sleep(1) +def ask_input(question, default='', numeric_only=False, valid_responses=[]): + question_out = question + if default: + question_out += f' (Default: {default})' + response = '' + while True: + response = input(question_out + ' ').strip() + if not response and default: + return default + if numeric_only: + if not response.isnumeric(): + print("Invalid value, try again.") + continue + response = int(response) + if valid_responses and response not in valid_responses: + print("Invalid value, try again.") + continue + break + return response def ask_yes_no(question, default=False) -> bool: """Ask a yes/no question via input() and return their answer.""" @@ -672,6 +710,19 @@ def dict_search(path, my_dict): c = c.get(p, {}) return c.get(parts[-1], None) +def dict_search_args(dict_object, *path): + # Traverse dictionary using variable arguments + # Added due to above function not allowing for '.' in the key names + # Example: dict_search_args(some_dict, 'key', 'subkey', 'subsubkey', ...) + if not isinstance(dict_object, dict) or not path: + return None + + for item in path: + if item not in dict_object: + return None + dict_object = dict_object[item] + return dict_object + def get_interface_config(interface): """ Returns the used encapsulation protocol for given interface. If interface does not exist, None is returned. @@ -682,6 +733,16 @@ def get_interface_config(interface): tmp = loads(cmd(f'ip -d -j link show {interface}'))[0] return tmp +def get_interface_address(interface): + """ Returns the used encapsulation protocol for given interface. + If interface does not exist, None is returned. + """ + if not os.path.exists(f'/sys/class/net/{interface}'): + return None + from json import loads + tmp = loads(cmd(f'ip -d -j addr show {interface}'))[0] + return tmp + def get_all_vrfs(): """ Return a dictionary of all system wide known VRF instances """ from json import loads @@ -694,3 +755,58 @@ def get_all_vrfs(): name = entry.pop('name') data[name] = entry return data + +def print_error(str='', end='\n'): + """ + Print `str` to stderr, terminated with `end`. + Used for warnings and out-of-band messages to avoid mangling precious + stdout output. + """ + sys.stderr.write(str) + 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 is_systemd_service_running(service): + """ Test is a specified systemd service is actually running. + Returns True if service is running, false otherwise. """ + tmp = run(f'systemctl is-active --quiet {service}') + return bool((tmp == 0)) |