live: rewrite image_restore

Integrates image restore command into native ogClient code. Further
reduces the need for external Bash scripts.

After a succesful image restore, OS configuration is still using
external Bash script "osConfigure/osConfigureCustom".

3 files changed, 99 insertions, 21 deletions

diff --git a/src/live/ogOperations.py b/src/live/ogOperations.py
index d6e2b6c..e99eb41 100644
--- a/src/live/ogOperations.py
+++ b/src/live/ogOperations.py
@@ -11,6 +11,7 @@ import logging
 import os
 import subprocess
 import shlex
+import shutil
 
 from subprocess import Popen, PIPE
 
@@ -25,8 +26,9 @@ from src.utils.net import ethtool
 from src.utils.menu import generate_menu
 from src.utils.fs import *
 from src.utils.probe import os_probe, cache_probe
-from src.utils.disk import get_disks
+from src.utils.disk import *
 from src.utils.cache import generate_cache_txt
+from src.utils.tiptorrent import *
 
 
 OG_SHELL = '/bin/bash'
@@ -117,6 +119,36 @@ class OgLiveOperations:
         with open(filename, 'w') as f:
             f.write(dig)
 
+    def _restore_image_unicast(self, repo, name, devpath):
+        image_path = f'/opt/opengnsys/images/{name}.img'
+        if ogChangeRepo(repo).returncode != 0:
+            logging.error('ogChangeRepo could not change repository to %s', repo)
+            raise ValueError(f'Error: Cannot change repository to {repo}')
+        self._restore_image(image_path, devpath)
+
+    def _restore_image_tiptorrent(self, repo, name, devpath):
+        image_path = f'/opt/opengnsys/cache/opt/opengnsys/images/{name}.img'
+        if (not os.path.exists(image_path) or
+            not tip_check_csum(repo, name)):
+            tip_client_get(repo, name)
+        self._restore_image(image_path, devpath)
+
+    def _restore_image(self, image_path, devpath):
+        cmd_lzop = shlex.split('lzop -dc -')
+        cmd_pc = shlex.split(f'partclone.restore -d0 -C -I -o {devpath}')
+        cmd_mbuffer = shlex.split('mbuffer -q -m 40M') if shutil.which('mbuffer') else None
+
+        with open(image_path, 'rb') as imgfile, \
+             open('/tmp/command.log', 'wb', 0) as logfile:
+            proc_lzop = subprocess.Popen(cmd_lzop,
+                                         stdout=subprocess.PIPE,
+                                         stdin=imgfile)
+            proc_pc = subprocess.Popen(cmd_pc,
+                                       stdin=proc_lzop.stdout,
+                                       stderr=logfile)
+            proc_lzop.stdout.close()
+            proc_pc.communicate()
+
     def poweroff(self):
         logging.info('Powering off client')
         if os.path.exists('/scripts/oginit'):
@@ -265,27 +297,30 @@ class OgLiveOperations:
         ctype = request.getType()
         profile = request.getProfile()
         cid = request.getId()
-        cmd = f'{ogClient.OG_PATH}interfaceAdm/RestaurarImagen {disk} {partition} ' \
-              f'{name} {repo} {ctype}'
+
+        partdev = get_partition_device(int(disk), int(partition))
 
         self._restartBrowser(self._url_log)
 
-        try:
-            ogRest.proc = subprocess.Popen([cmd],
-                               stdout=subprocess.PIPE,
-                               shell=True,
-                               executable=OG_SHELL)
-            (output, error) = ogRest.proc.communicate()
-            if (ogRest.proc.returncode):
-                raise Exception
-        except:
-            logging.error('Exception when running image restore subprocess')
-            raise ValueError('Error: Incorrect command value')
+        logging.debug('Image restore params:')
+        logging.debug(f'\tname: {name}')
+        logging.debug(f'\trepo: {repo}')
+        logging.debug(f'\tprofile: {profile}')
+        logging.debug(f'\tctype: {ctype}')
+
+        if shutil.which('restoreImageCustom'):
+            restoreImageCustom(repo, name, disk, partition, ctype)
+        elif 'UNICAST' in ctype:
+            self._restore_image_unicast(repo, name, partdev)
+        elif ctype == 'TIPTORRENT':
+            self._restore_image_tiptorrent(repo, name, partdev)
+
+        output = configureOs(disk, partition)
 
         self.refresh(ogRest)
 
         logging.info('Image restore command OK')
-        return output.decode('utf-8')
+        return output
 
     def image_create(self, path, request, ogRest):
         disk = int(request.getDisk())
diff --git a/src/utils/legacy.py b/src/utils/legacy.py
index 8f39fe2..aca8abe 100644
--- a/src/utils/legacy.py
+++ b/src/utils/legacy.py
@@ -3,6 +3,7 @@ import logging
 import os
 import subprocess
 import shlex
+import shutil
 
 from subprocess import PIPE
 
@@ -56,6 +57,51 @@ def ogChangeRepo(ip):
                           shell=True)
 
 
+def restoreImageCustom(repo_ip, image_name, disk, partition, method):
+    """
+    """
+    if not shutil.which('restoreImageCustom'):
+        logging.error('Invalid restoreImageCustom invocation')
+        raise ValueError('Error: restoreImageCustom not found')
+
+    if ogChangeRepo(repo).returncode != 0:
+        logging.error('ogChangeRepo could not change repository to %s', repo)
+        raise ValueError(f'Error: Cannot change repository to {repo}')
+
+    cmd = f'restoreImageCustom {repo_ip} {image_name} {disk} {partition} {method}'
+    with open('/tmp/command.log', 'wb', 0) as logfile:
+        try:
+            proc = subprocess.run(cmd,
+                                  stdout=logfile,
+                                  encoding='utf-8',
+                                  shell=True)
+        except:
+            logging.error('Exception when running restoreImageCustom subprocess')
+            raise ValueError('Error: Incorrect command value')
+    return proc.returncode
+
+
+def configureOs(disk, partition):
+    """
+    """
+    if shutil.which('configureOsCustom'):
+        cmd_configure = f"configureOsCustom {disk} {partition}"
+    else:
+        cmd_configure = f"configureOs {disk} {partition}"
+
+    try:
+        proc = subprocess.run(cmd_configure,
+                              stdout=PIPE,
+                              encoding='utf-8',
+                              shell=True)
+        out = proc.stdout
+    except:
+        logging.error('Exception when running configureOs subprocess')
+        raise ValueError('Error: Incorrect command value')
+
+    return out
+
+
 def ogCopyEfiBootLoader(disk, partition):
     cmd = f'ogCopyEfiBootLoader {disk} {partition}'
     try:
diff --git a/src/utils/tiptorrent.py b/src/utils/tiptorrent.py
index 796a665..62feb43 100644
--- a/src/utils/tiptorrent.py
+++ b/src/utils/tiptorrent.py
@@ -71,16 +71,13 @@ def tip_client_get(tip_addr, image_name):
 
     try:
         proc = subprocess.Popen(shlex.split(cmd),
-                                stdout=subprocess.PIPE,
-                                stderr=subprocess.PIPE,
+                                stdout=logfile,
                                 cwd='/opt/opengnsys/cache/opt/opengnsys/images/')
-        out, err = proc.communicate()
+        proc.communicate()
     except:
-        logging.error('Exception when running software inventory subprocess')
+        logging.error('Exception when running tiptorrent client GET subprocess')
         raise ValueError('Error: Incorrect command value')
     finally:
-        logging.debug(f'tip_client_get out: {out}')
-        logging.debug(f'tip_client_get err: {err}')
         logfile.close()
 
     if proc.returncode != 0: