From 50197c6a12c0928dab80eba1b31d0061bb630be3 Mon Sep 17 00:00:00 2001
From: AlanCoding <arominge@redhat.com>
Date: Sat, 18 Apr 2020 23:17:59 -0400
Subject: [PATCH] Handle custom script options of hostnames and destination
 variable

---
 awx/main/models/inventory.py                  | 29 +++++++++++++++----
 .../inventory/plugins/ec2/files/aws_ec2.yml   |  4 +--
 .../vmware/files/vmware_vm_inventory.yml      |  2 +-
 .../scripts/ec2/files/file_reference          |  3 +-
 .../test_inventory_source_injectors.py        |  4 ++-
 5 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/awx/main/models/inventory.py b/awx/main/models/inventory.py
index 3308c84b8b..a32a24da44 100644
--- a/awx/main/models/inventory.py
+++ b/awx/main/models/inventory.py
@@ -2053,18 +2053,35 @@ class ec2(PluginFileInjector):
         compose_dict.update(self._compat_compose_vars())
         # plugin provides "aws_ec2", but not this which the script gave
         ret['groups'] = {'ec2': True}
-        # public_ip as hostname is non-default plugin behavior, script behavior
-        ret['hostnames'] = [
-            'network-interface.addresses.association.public-ip',
-            'dns-name',
-            'private-dns-name'
-        ]
+        if source_vars.get('hostname_variable') is not None:
+            hnames = []
+            for expr in source_vars.get('hostname_variable').split(','):
+                if expr == 'public_dns_name':
+                    hnames.append('dns-name')
+                elif not expr.startswith('tag:') and '_' in expr:
+                    hnames.append(expr.replace('_', '-'))
+                else:
+                    hnames.append(expr)
+            ret['hostnames'] = hnames
+        else:
+            # public_ip as hostname is non-default plugin behavior, script behavior
+            ret['hostnames'] = [
+                'network-interface.addresses.association.public-ip',
+                'dns-name',
+                'private-dns-name'
+            ]
         # The script returned only running state by default, the plugin does not
         # https://docs.aws.amazon.com/cli/latest/reference/ec2/describe-instances.html#options
         # options: pending | running | shutting-down | terminated | stopping | stopped
         inst_filters['instance-state-name'] = ['running']
         # end compatibility content
 
+        if source_vars.get('destination_variable') or source_vars.get('vpc_destination_variable'):
+            for fd in ('destination_variable', 'vpc_destination_variable'):
+                if source_vars.get(fd):
+                    compose_dict['ansible_host'] = source_vars.get(fd)
+                    break
+
         if compose_dict:
             ret['compose'] = compose_dict
 
diff --git a/awx/main/tests/data/inventory/plugins/ec2/files/aws_ec2.yml b/awx/main/tests/data/inventory/plugins/ec2/files/aws_ec2.yml
index fc7f9c2421..8984d4cb56 100644
--- a/awx/main/tests/data/inventory/plugins/ec2/files/aws_ec2.yml
+++ b/awx/main/tests/data/inventory/plugins/ec2/files/aws_ec2.yml
@@ -1,6 +1,6 @@
 boto_profile: /tmp/my_boto_stuff
 compose:
-  ansible_host: public_ip_address
+  ansible_host: public_dns_name
   ec2_account_id: owner_id
   ec2_ami_launch_index: ami_launch_index | string
   ec2_architecture: architecture
@@ -50,9 +50,7 @@ filters:
 groups:
   ec2: true
 hostnames:
-- network-interface.addresses.association.public-ip
 - dns-name
-- private-dns-name
 iam_role_arn: arn:aws:iam::123456789012:role/test-role
 keyed_groups:
 - key: placement.availability_zone
diff --git a/awx/main/tests/data/inventory/plugins/vmware/files/vmware_vm_inventory.yml b/awx/main/tests/data/inventory/plugins/vmware/files/vmware_vm_inventory.yml
index 3bc5269dd8..18b96b9e78 100644
--- a/awx/main/tests/data/inventory/plugins/vmware/files/vmware_vm_inventory.yml
+++ b/awx/main/tests/data/inventory/plugins/vmware/files/vmware_vm_inventory.yml
@@ -39,5 +39,5 @@ properties:
 - triggeredalarmstate
 - value
 property_name_format: lower_case
-strict: true
+strict: false
 with_nested_properties: true
diff --git a/awx/main/tests/data/inventory/scripts/ec2/files/file_reference b/awx/main/tests/data/inventory/scripts/ec2/files/file_reference
index 564e631ba0..4ffb85656f 100644
--- a/awx/main/tests/data/inventory/scripts/ec2/files/file_reference
+++ b/awx/main/tests/data/inventory/scripts/ec2/files/file_reference
@@ -2,9 +2,10 @@
 base_source_var = value_of_var
 boto_profile = /tmp/my_boto_stuff
 iam_role_arn = arn:aws:iam::123456789012:role/test-role
+hostname_variable = public_dns_name
+destination_variable = public_dns_name
 regions = us-east-2,ap-south-1
 regions_exclude = us-gov-west-1,cn-north-1
-destination_variable = public_dns_name
 vpc_destination_variable = ip_address
 route53 = False
 all_instances = True
diff --git a/awx/main/tests/functional/test_inventory_source_injectors.py b/awx/main/tests/functional/test_inventory_source_injectors.py
index e776753274..50510e3c16 100644
--- a/awx/main/tests/functional/test_inventory_source_injectors.py
+++ b/awx/main/tests/functional/test_inventory_source_injectors.py
@@ -39,7 +39,9 @@ TEST_SOURCE_FIELDS = {
 INI_TEST_VARS = {
     'ec2': {
         'boto_profile': '/tmp/my_boto_stuff',
-        'iam_role_arn': 'arn:aws:iam::123456789012:role/test-role'
+        'iam_role_arn': 'arn:aws:iam::123456789012:role/test-role',
+        'hostname_variable': 'public_dns_name',
+        'destination_variable': 'public_dns_name'
     },
     'gce': {},
     'openstack': {