frr: T1514: refactor restart script and drop duplicated code

1 files changed, 57 insertions, 74 deletions

diff --git a/src/op_mode/restart_frr.py b/src/op_mode/restart_frr.py
index 0b2322478..109c8dd7b 100755
--- a/src/op_mode/restart_frr.py
+++ b/src/op_mode/restart_frr.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 #
-# Copyright (C) 2019 VyOS maintainers and contributors
+# Copyright (C) 2019-2021 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
@@ -13,16 +13,19 @@
 #
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
-#
 
-import sys
+import os
 import argparse
 import logging
-from logging.handlers import SysLogHandler
-from pathlib import Path
 import psutil
 
+from logging.handlers import SysLogHandler
+from shutil import rmtree
+
 from vyos.util import call
+from vyos.util import ask_yes_no
+from vyos.util import process_named_running
+from vyos.util import makedir
 
 # some default values
 watchfrr = '/usr/lib/frr/watchfrr.sh'
@@ -40,40 +43,45 @@ logger.setLevel(logging.INFO)
 def _check_safety():
     try:
         # print warning
-        answer = input("WARNING: This is a potentially unsafe function! You may lose the connection to the router or active configuration after running this command. Use it at your own risk! Continue? [y/N]: ")
-        if not answer.lower() == "y":
-            logger.error("User aborted command")
+        if not ask_yes_no('WARNING: This is a potentially unsafe function!\n' \
+                          'You may lose the connection to the router or active configuration after\n' \
+                          'running this command. Use it at your own risk!\n\n'
+                          'Continue?'):
             return False
 
         # check if another restart process already running
         if len([process for process in psutil.process_iter(attrs=['pid', 'name', 'cmdline']) if 'python' in process.info['name'] and 'restart_frr.py' in process.info['cmdline'][1]]) > 1:
-            logger.error("Another restart_frr.py already running")
-            answer = input("Another restart_frr.py process is already running. It is unsafe to continue. Do you want to process anyway? [y/N]: ")
-            if not answer.lower() == "y":
+            message = 'Another restart_frr.py process is already running!'
+            logger.error(message)
+            if not ask_yes_no(f'\n{message} It is unsafe to continue.\n\n' \
+                              'Do you want to process anyway?'):
                 return False
 
         # check if watchfrr.sh is running
-        for process in psutil.process_iter(attrs=['pid', 'name', 'cmdline']):
-            if 'bash' in process.info['name'] and watchfrr in process.info['cmdline']:
-                logger.error("Another {} already running".format(watchfrr))
-                answer = input("Another {} process is already running. It is unsafe to continue. Do you want to process anyway? [y/N]: ".format(watchfrr))
-                if not answer.lower() == "y":
-                    return False
+        tmp = os.path.basename(watchfrr)
+        if process_named_running(tmp):
+            message = f'Another {tmp} process is already running.'
+            logger.error(message)
+            if not ask_yes_no(f'{message} It is unsafe to continue.\n\n' \
+                              'Do you want to process anyway?'):
+                return False
 
         # check if vtysh is running
-        for process in psutil.process_iter(attrs=['pid', 'name', 'cmdline']):
-            if 'vtysh' in process.info['name']:
-                logger.error("The vtysh is running by another task")
-                answer = input("The vtysh is running by another task. It is unsafe to continue. Do you want to process anyway? [y/N]: ")
-                if not answer.lower() == "y":
-                    return False
+        if process_named_running('vtysh'):
+            message = 'vtysh process is executed by another task.'
+            logger.error(message)
+            if not ask_yes_no(f'{message} It is unsafe to continue.\n\n' \
+                              'Do you want to process anyway?'):
+                return False
 
         # check if temporary directory exists
-        if Path(frrconfig_tmp).exists():
-            logger.error("The temporary directory \"{}\" already exists".format(frrconfig_tmp))
-            answer = input("The temporary directory \"{}\" already exists. It is unsafe to continue. Do you want to process anyway? [y/N]: ".format(frrconfig_tmp))
-            if not answer.lower() == "y":
+        if os.path.exists(frrconfig_tmp):
+            message = f'Temporary directory "{frrconfig_tmp}" already exists!'
+            logger.error(message)
+            if not ask_yes_no(f'{message} It is unsafe to continue.\n\n' \
+                              'Do you want to process anyway?'):
                 return False
+
     except:
         logger.error("Something goes wrong in _check_safety()")
         return False
@@ -84,72 +92,47 @@ def _check_safety():
 # write active config to file
 def _write_config():
     # create temporary directory
-    Path(frrconfig_tmp).mkdir(parents=False, exist_ok=True)
+    makedir(frrconfig_tmp)
     # save frr.conf to it
-    command = "{} -n -w --config_dir {} 2> /dev/null".format(vtysh, frrconfig_tmp)
+    command = f'{vtysh} -n -w --config_dir {frrconfig_tmp} 2> /dev/null'
     return_code = call(command)
-    if not return_code == 0:
-        logger.error("Failed to save active config: \"{}\" returned exit code: {}".format(command, return_code))
+    if return_code != 0:
+        logger.error(f'Failed to save active config: "{command}" returned exit code: {return_code}')
         return False
-    logger.info("Active config saved to {}".format(frrconfig_tmp))
+    logger.info(f'Active config saved to {frrconfig_tmp}')
     return True
 
 # clear and remove temporary directory
 def _cleanup():
-    tmpdir = Path(frrconfig_tmp)
-    try:
-        if tmpdir.exists():
-            for file in tmpdir.iterdir():
-                file.unlink()
-            tmpdir.rmdir()
-    except:
-        logger.error("Failed to remove temporary directory {}".format(frrconfig_tmp))
-        print("Failed to remove temporary directory {}".format(frrconfig_tmp))
-
-# check if daemon is running
-def _daemon_check(daemon):
-    command = "{} print_status {}".format(watchfrr, daemon)
-    return_code = call(command)
-    if not return_code == 0:
-        logger.error("Daemon \"{}\" is not running".format(daemon))
-        return False
-
-    # return True if all checks were passed
-    return True
+    if os.path.isdir(frrconfig_tmp):
+        rmtree(frrconfig_tmp)
 
 # restart daemon
 def _daemon_restart(daemon):
-    command = "{} restart {}".format(watchfrr, daemon)
+    command = f'{watchfrr} restart {daemon}'
     return_code = call(command)
     if not return_code == 0:
-        logger.error("Failed to restart daemon \"{}\"".format(daemon))
+        logger.error(f'Failed to restart daemon "{daemon}"!')
         return False
 
     # return True if restarted successfully
-    logger.info("Daemon \"{}\" restarted".format(daemon))
+    logger.info(f'Daemon "{daemon}" restarted!')
     return True
 
 # reload old config
 def _reload_config(daemon):
     if daemon != '':
-        command = "{} -n -b --config_dir {} -d {} 2> /dev/null".format(vtysh, frrconfig_tmp, daemon)
+        command = f'{vtysh} -n -b --config_dir {frrconfig_tmp} -d {daemon} 2> /dev/null'
     else:
-        command = "{} -n -b --config_dir {} 2> /dev/null".format(vtysh, frrconfig_tmp)
+        command = f'{vtysh} -n -b --config_dir {frrconfig_tmp} 2> /dev/null'
 
     return_code = call(command)
     if not return_code == 0:
-        logger.error("Failed to reinstall configuration")
+        logger.error('Failed to re-install configuration!')
         return False
 
     # return True if restarted successfully
-    logger.info("Configuration reinstalled successfully")
-    return True
-
-# check all daemons if they are running
-def _check_args_daemon(daemons):
-    for daemon in daemons:
-        if not _daemon_check(daemon):
-            return False
+    logger.info('Configuration re-installed successfully!')
     return True
 
 # define program arguments
@@ -159,19 +142,18 @@ cmd_args_parser.add_argument('--daemon', choices=['bfdd', 'bgpd', 'ospfd', 'ospf
 # parse arguments
 cmd_args = cmd_args_parser.parse_args()
 
-
 # main logic
 # restart daemon
 if cmd_args.action == 'restart':
     # check if it is safe to restart FRR
     if not _check_safety():
         print("\nOne of the safety checks was failed or user aborted command. Exiting.")
-        sys.exit(1)
+        exit(1)
 
     if not _write_config():
         print("Failed to save active config")
         _cleanup()
-        sys.exit(1)
+        exit(1)
 
     # a little trick to make further commands more clear
     if not cmd_args.daemon:
@@ -179,19 +161,20 @@ if cmd_args.action == 'restart':
 
     # check all daemons if they are running
     if cmd_args.daemon != ['']:
-        if not _check_args_daemon(cmd_args.daemon):
-            print("Warning: some of listed daemons are not running")
+        for daemon in cmd_args.daemon:
+            if not process_named_running(daemon):
+                print('WARNING: some of listed daemons are not running!')
 
     # run command to restart daemon
     for daemon in cmd_args.daemon:
         if not _daemon_restart(daemon):
-            print("Failed to restart daemon: {}".format(daemon))
+            print('Failed to restart daemon: {daemon}')
             _cleanup()
-            sys.exit(1)
+            exit(1)
         # reinstall old configuration
         _reload_config(daemon)
 
     # cleanup after all actions
     _cleanup()
 
-sys.exit(0)
+exit(0)