views: remove unsupported partition types from image operations

Prevent unexpected behaviour in image operations caused by user
error.

Remove partitions of type 'EMPTY', 'LINUX-SWAP', 'CACHE' and  'EFI'
from the image/create, image/restore and image/update forms.

Remove check for the existence of these partitions in
action_image_restore as they are no longer available in the form.

1 files changed, 31 insertions, 12 deletions

diff --git a/ogcp/views.py b/ogcp/views.py
index 59e138c..92dbf32 100644
--- a/ogcp/views.py
+++ b/ogcp/views.py
@@ -77,6 +77,9 @@ PART_TYPE_CODES = {
     65535: 'UNKNOWN'
 }
 
+def get_invalid_image_partition_types():
+    return ['EMPTY', 'LINUX-SWAP', 'CACHE', 'EFI']
+
 PART_SCHEME_CODES = {
     0: 'EMPTY',
     1: 'MSDOS',
@@ -783,12 +786,6 @@ def action_image_restore():
         ips = form.ips.data.split(' ')
         disk, partition, part_code, part_size, has_cache = form.partition.data.split(' ')
 
-        part_type = PART_TYPE_CODES.get(int(part_code), 'UNKNOWN')
-        invalid_part_types = ['EMPTY', 'LINUX-SWAP', 'CACHE', 'EFI', 'WIN-RECOV']
-        if  part_type in invalid_part_types:
-            flash(_(f'Cannot restore image on partition type {part_type}'), category='error')
-            return redirect(url_for('commands'))
-
         requires_cache = form.method.data == 'TIPTORRENT' or form.method.data == 'UNICAST'
         if has_cache == 'False' and requires_cache:
             flash(_(f'Cannot restore image using {form.method.data} on a client without cache'), category='error')
@@ -902,12 +899,20 @@ def action_image_restore():
                 flash(_(f'Computers have different partition setup'), category='error')
                 return redirect(url_for('commands'))
 
-        form.partition.choices = [
-            (f"{disk_id} {part_id} {part_code} {part_size} {has_cache}",
-             f"Disk {disk_id} | Partition {part_id} "
-             f"| {PART_TYPE_CODES.get(part_code, 'UNKNOWN')} "
-             f"{FS_CODES.get(filesystem, 'UNKNOWN')}")
-            for disk_id, part_id, part_code, filesystem, part_size in part_choices ]
+        invalid_part_types = get_invalid_image_partition_types()
+
+        for disk_id, part_id, part_code, filesystem, part_size in part_choices:
+            part_type = PART_TYPE_CODES.get(int(part_code), 'UNKNOWN')
+
+            if  part_type in invalid_part_types:
+                continue
+
+            form.partition.choices.append(
+                (f"{disk_id} {part_id} {part_code} {part_size} {has_cache}",
+                f"Disk {disk_id} | Partition {part_id} "
+                f"| {PART_TYPE_CODES.get(part_code, 'UNKNOWN')} "
+                f"{FS_CODES.get(filesystem, 'UNKNOWN')}")
+            )
         scopes, clients = get_scopes(set(ips))
         selected_clients = list(get_selected_clients(scopes['scope']).items())
 
@@ -2012,7 +2017,14 @@ def action_image_create():
         if r.status_code != requests.codes.ok:
             return ogserver_error('commands')
 
+        invalid_part_types = get_invalid_image_partition_types()
+
         for part in r.json()['partitions'][1:]:
+            part_type = PART_TYPE_CODES.get(int(part.get('code')), 'UNKNOWN')
+
+            if part_type in invalid_part_types:
+                continue
+
             form.os.choices.append(
                 (f"{part.get('disk')} {part.get('partition')} {part.get('code')}",
                  f"Disk {part.get('disk')} | Partition {part.get('partition')} "
@@ -2119,7 +2131,14 @@ def action_image_update():
     if r.status_code != requests.codes.ok:
         return ogserver_error('commands')
 
+    invalid_part_types = get_invalid_image_partition_types()
+
     for part in r.json()['partitions'][1:]:
+        part_type = PART_TYPE_CODES.get(int(part.get('code')), 'UNKNOWN')
+
+        if part_type in invalid_part_types:
+            continue
+
         form.os.choices.append(
             (f"{part.get('disk')} {part.get('partition')} {part.get('code')}",
              f"Disk {part.get('disk')} | Partition {part.get('partition')} "