load-config: T2138: support loading gzipped (local) config files

1 files changed, 30 insertions, 16 deletions

diff --git a/src/helpers/vyos-load-config.py b/src/helpers/vyos-load-config.py
index c2da1bb11..e579e81b2 100755
--- a/src/helpers/vyos-load-config.py
+++ b/src/helpers/vyos-load-config.py
@@ -23,7 +23,9 @@ Example: load https://somewhere.net/some.config
          load /tmp/some.config
 """
 
+import os
 import sys
+import gzip
 import tempfile
 import vyos.defaults
 import vyos.remote
@@ -37,35 +39,47 @@ class LoadConfig(ConfigSourceSession):
     def load_config(self, path):
         return self._run(['/bin/cli-shell-api','loadFile',path])
 
-
 file_name = sys.argv[1] if len(sys.argv) > 1 else 'config.boot'
 configdir = vyos.defaults.directories['config']
 protocols = ['scp', 'sftp', 'http', 'https', 'ftp', 'tftp']
 
+def get_local_config(filename):
+    if os.path.isfile(filename):
+        fname = filename
+    elif os.path.isfile(os.path.join(configdir, filename)):
+        fname = os.path.join(configdir, filename)
+    else:
+        sys.exit(f"No such file '{filename}'")
+
+    if fname.endswith('.gz'):
+        with gzip.open(fname, 'rb') as f:
+            try:
+                config_str = f.read().decode()
+            except OSError as e:
+                sys.exit(e)
+    else:
+        with open(fname, 'r') as f:
+            try:
+                config_str = f.read()
+            except OSError as e:
+                sys.exit(e)
+
+    return config_str
 
 if any(x in file_name for x in protocols):
-    config_file = vyos.remote.get_remote_config(file_name)
-    if not config_file:
-        sys.exit("No config file by that name.")
+    config_string = vyos.remote.get_remote_config(file_name)
+    if not config_string:
+        sys.exit(f"No such config file at '{file_name}'")
 else:
-    canonical_path = '{0}/{1}'.format(configdir, file_name)
-    try:
-        with open(canonical_path, 'r') as f:
-            config_file = f.read()
-    except OSError as err1:
-        try:
-            with open(file_name, 'r') as f:
-                config_file = f.read()
-        except OSError as err2:
-            sys.exit('{0}\n{1}'.format(err1, err2))
+    config_string = get_local_config(file_name)
 
 config = LoadConfig()
 
-print("Loading configuration from '{}'".format(file_name))
+print(f"Loading configuration from '{file_name}'")
 
 with tempfile.NamedTemporaryFile() as fp:
     with open(fp.name, 'w') as fd:
-        fd.write(config_file)
+        fd.write(config_string)
 
     virtual_migration = VirtualMigrator(fp.name)
     try: