Limit restore image to images of assigned repo

Allow to restore only to images that are in the repo the client
is assigned to.

1 files changed, 35 insertions, 4 deletions

diff --git a/ogcp/views.py b/ogcp/views.py
index 8494d4f..a10c9c4 100644
--- a/ogcp/views.py
+++ b/ogcp/views.py
@@ -612,6 +612,30 @@ def search_image(images_list, image_id):
             return image
     return False
 
+def get_images_grouped_by_repos_from_server(server):
+    r = server.get('/images')
+    images = r.json()['images']
+    repos={}
+
+    for image in images:
+        repo_id=image['repo_id']
+        if repo_id not in repos:
+            repos[repo_id] = [image]
+        else:
+            repos[repo_id].append(image)
+    return repos
+
+def get_clients_repo(server, ips):
+    repo_id=None
+    for ip in ips:
+        r = server.get('/client/info', payload={'client': [ip]})
+        repo_id_aux = r.json()['repo_id']
+        if repo_id is None:
+            repo_id = repo_id_aux
+        elif repo_id_aux != repo_id:
+            return None
+    return repo_id
+
 @app.route('/action/image/restore', methods=['GET', 'POST'])
 @login_required
 def action_image_restore():
@@ -649,13 +673,20 @@ def action_image_restore():
             return redirect(url_for('commands'))
         form.ips.data = ' '.join(ips)
 
-        part_choices = []
-
         server = get_server_from_clients(ips)
-        r = server.get('/images')
-        for image in r.json()['images']:
+
+        repo_id = get_clients_repo(server, ips)
+        if repo_id is None:
+            flash(_(f'Computers have different repos assigned'), category='error')
+            return redirect(url_for('commands'))
+        images = get_images_grouped_by_repos_from_server(server)
+        if repo_id not in images:
+            flash(_(f'Computer(s) assigned to a repo with no images'), category='error')
+            return redirect(url_for('commands'))
+        for image in images[repo_id]:
             form.image.choices.append((image['id'], image['name']))
 
+        part_choices = []
         for ip in ips:
             r = server.get('/client/setup', payload={'client': [ip]})
             if r.status_code == requests.codes.ok: