From ccd1de8280fa2d560b35117a61f15f67267f7fea Mon Sep 17 00:00:00 2001
From: "Z. Cliffe Schreuders" <code.cliffe@schreuders.org>
Date: Thu, 25 Oct 2018 16:24:47 +0100
Subject: [PATCH] affinity groups

---
 lib/helpers/ovirt.rb          | 16 ++++++-
 lib/helpers/ovirt_affinity.py | 81 +++++++++++++++++++++++++++++++++++
 secgen.rb                     |  8 ++++
 3 files changed, 104 insertions(+), 1 deletion(-)
 create mode 100644 lib/helpers/ovirt_affinity.py

diff --git a/lib/helpers/ovirt.rb b/lib/helpers/ovirt.rb
index f58c10849..15de743e2 100644
--- a/lib/helpers/ovirt.rb
+++ b/lib/helpers/ovirt.rb
@@ -200,7 +200,6 @@ class OVirtFunctions
       end
     end
 
-
     vms.each do |vm_list|
       vm_list.each do |vm|
         Print.std "  Assigning network to: #{vm.name}"
@@ -232,6 +231,21 @@ class OVirtFunctions
     end
   end
 
+  def self.assign_affinity_group(options, scenario_path, vm_names)
+    vms = []
+    ovirt_vm_names = build_ovirt_names(scenario_path, options[:prefix], vm_names)
+    ovirt_vm_names.each do |vm_name|
+      # python affinity group
+      if system "python #{ROOT_DIR}/lib/helpers/ovirt_affinity.py #{options[:ovirtaffinitygroup]} #{vm_name} #{options[:ovirturl]} #{options[:ovirtuser]}
+                   #{options[:ovirtpass]}"
+        Print.std "Affinity group assigned"
+      else
+        Print.err "Failed to assign affinity group"
+        exit 1
+      end
+    end
+  end
+
   def self.assign_permissions(options, scenario_path, vm_names)
     ovirt_connection = get_ovirt_connection(options)
     username = options[:prefix].chomp
diff --git a/lib/helpers/ovirt_affinity.py b/lib/helpers/ovirt_affinity.py
new file mode 100644
index 000000000..d11b69346
--- /dev/null
+++ b/lib/helpers/ovirt_affinity.py
@@ -0,0 +1,81 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+# The ruby ovirt sdk module doesn't access affinity groups correctly
+
+# by Paul Staniforth
+# and Z. Cliffe Schreuders
+
+import logging
+import getpass
+
+import ovirtsdk4 as sdk
+import ovirtsdk4.types as types
+
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument("affinitygroup")
+parser.add_argument("vm_name_search")
+parser.add_argument("ovirt_url")
+parser.add_argument("ovirt_username")
+parser.add_argument("ovirt_password")
+args = parser.parse_args()
+print(args.vm_name)
+print(args.affinitygroup)
+
+
+# logging.basicConfig(level=logging.DEBUG, filename='example.log')
+
+# Create the connection to the server:
+connection = sdk.Connection(
+    url=args.ovirt_url,
+    username=args.ovirt_username,
+    password=args.ovirt_password,
+    debug=True,
+    log=logging.getLogger(),
+)
+
+# Locate the clusters service and use it to find the cluster
+clusters_service = connection.system_service().clusters_service()
+cluster = clusters_service.list(search='name=default')[0]
+
+cluster_service = clusters_service.cluster_service(cluster.id)
+cluster_affinitygroups_service = cluster_service.affinity_groups_service()
+
+cluster_service = clusters_service.cluster_service(cluster.id)
+cluster_affinitygroups_service = cluster_service.affinity_groups_service()
+
+# cluster_affinitygroups_service.add(
+#      types.AffinityGroup(
+#         name='new_affinity_label10',
+#         description='software defined',
+#         vms_rule=types.AffinityRule(
+#              enabled=True,
+#              positive=True,
+#              enforcing=True,
+#         ),
+#      ),
+# )
+
+# Get the reference to the "vms" service:
+vms_service = connection.system_service().vms_service()
+
+# Find the virtual machine:
+vms = vms_service.list(search='name=' + args.vm_name_search)
+
+affinitygroups = cluster_affinitygroups_service.list()
+
+for affinitygroup in affinitygroups:
+    print ("Affinity_Group: %s Affinity_Group ID: %s Description: %s Comment: %s"%(affinitygroup.name,affinitygroup.id,affinitygroup.description,affinitygroup.comment))
+    if affinitygroup.name == args.affinitygroup:
+        group_service = cluster_affinitygroups_service.group_service(affinitygroup.id)
+        group_vms_service = group_service.vms_service()
+        for vm in vms:
+            group_vms_service.add(
+                vm=types.Vm(
+                    id=vm.id,
+                )
+            )
+
+# Close the connection to the server:
+connection.close()
diff --git a/secgen.rb b/secgen.rb
index 1dd45549a..5bb357bcc 100644
--- a/secgen.rb
+++ b/secgen.rb
@@ -193,6 +193,10 @@ def ovirt_post_build(options, scenario, project_dir)
     Print.info 'Assigning network(s) of VM(s)'
     OVirtFunctions::assign_networks(options, scenario, get_vm_names(scenario))
   end
+  if options[:ovirtaffinitygroup]
+    Print.info 'Assigning affinity group of VM(s)'
+    OVirtFunctions::assign_affinity_group(options, scenario, get_vm_names(scenario))
+  end
   if options[:snapshot]
     Print.info 'Creating a snapshot of VM(s)'
     if OVirtFunctions::provider_ovirt?(options)
@@ -363,6 +367,7 @@ opts = GetoptLong.new(
     ['--ovirtauthz', GetoptLong::REQUIRED_ARGUMENT],
     ['--ovirt-cluster', GetoptLong::REQUIRED_ARGUMENT],
     ['--ovirt-network', GetoptLong::REQUIRED_ARGUMENT],
+    ['--ovirt-affinity-group', GetoptLong::REQUIRED_ARGUMENT],
     ['--snapshot', GetoptLong::NO_ARGUMENT],
 )
 
@@ -451,6 +456,9 @@ opts.each do |opt, arg|
     when '--ovirt-network'
       Print.info "Ovirt Network Name : #{arg}"
       options[:ovirtnetwork] = arg
+    when '--ovirt-affinity-group'
+      Print.info "Ovirt Affinity Group : #{arg}"
+      options[:ovirtaffinitygroup] = arg
     when '--snapshot'
       Print.info "Taking snapshots when VMs are created"
       options[:snapshot] = true