From 4d39b6af51769788b56abfc663ba33d079a8dcc9 Mon Sep 17 00:00:00 2001
From: Darko Poljak <darko.poljak@gmail.com>
Date: Thu, 15 Feb 2018 20:33:36 +0100
Subject: [PATCH] Add -4 and -6 params to force IPv4, IPv6 addresses
 respectively.

---
 cdist/argparse.py       | 10 ++++++++++
 cdist/config.py         | 28 +++++++++++++++++++++++++---
 cdist/util/ipaddr.py    |  8 ++++----
 docs/src/man1/cdist.rst | 18 ++++++++++++++----
 4 files changed, 53 insertions(+), 11 deletions(-)

diff --git a/cdist/argparse.py b/cdist/argparse.py
index 2ec28121..a9f3ad9f 100644
--- a/cdist/argparse.py
+++ b/cdist/argparse.py
@@ -224,6 +224,16 @@ def get_parsers():
 
     # Config
     parser['config_args'] = argparse.ArgumentParser(add_help=False)
+    parser['config_args'].add_argument(
+           '-4', '--force-ipv4',
+           help=('Force to use IPv4 addresses only. No influence for custom'
+                 ' remote commands.'),
+           action='store_const', dest='force_ipv', const=4)
+    parser['config_args'].add_argument(
+           '-6', '--force-ipv6',
+           help=('Force to use IPv6 addresses only. No influence for custom'
+                 ' remote commands.'),
+           action='store_const', dest='force_ipv', const=6)
     parser['config_args'].add_argument(
              '-A', '--all-tagged',
              help=('Use all hosts present in tags db. Currently in beta.'),
diff --git a/cdist/config.py b/cdist/config.py
index be39de22..63566594 100644
--- a/cdist/config.py
+++ b/cdist/config.py
@@ -32,6 +32,7 @@ import multiprocessing
 from cdist.mputil import mp_pool_run, mp_sig_handler
 import atexit
 import shutil
+import socket
 import cdist
 import cdist.hostsource
 import cdist.exec.local
@@ -110,6 +111,13 @@ class Config(object):
         args.remote_exec_pattern = None
         args.remote_copy_pattern = None
 
+        # Determine forcing IPv4/IPv6 options if any, only for
+        # default remote commands.
+        if args.force_ipv:
+            force_addr_opt = " -{}".format(args.force_ipv)
+        else:
+            force_addr_opt = ""
+
         args_dict = vars(args)
         # if remote-exec and/or remote-copy args are None then user
         # didn't specify command line options nor env vars:
@@ -118,9 +126,11 @@ class Config(object):
                 args_dict['remote_exec'] is None):
             mux_opts = inspect_ssh_mux_opts()
             if args_dict['remote_exec'] is None:
-                args.remote_exec_pattern = cdist.REMOTE_EXEC + mux_opts
+                args.remote_exec_pattern = (cdist.REMOTE_EXEC +
+                                            force_addr_opt + mux_opts)
             if args_dict['remote_copy'] is None:
-                args.remote_copy_pattern = cdist.REMOTE_COPY + mux_opts
+                args.remote_copy_pattern = (cdist.REMOTE_COPY +
+                                            force_addr_opt + mux_opts)
             if mux_opts:
                 cleanup_pattern = cdist.REMOTE_CMDS_CLEANUP_PATTERN
             else:
@@ -313,6 +323,16 @@ class Config(object):
             remote_cmds_cleanup = ""
         return (remote_exec, remote_copy, remote_cmds_cleanup, )
 
+    @staticmethod
+    def _address_family(args):
+        if args.force_ipv == 4:
+            family = socket.AF_INET
+        elif args.force_ipv == 6:
+            family = socket.AF_INET6
+        else:
+            family = 0
+        return family
+
     @classmethod
     def onehost(cls, host, host_tags, host_base_path, host_dir_name, args,
                 parallel, configuration, remove_remote_files_dirs=False):
@@ -331,7 +351,9 @@ class Config(object):
             log.debug("remote_copy for host \"{}\": {}".format(
                 host, remote_copy))
 
-            target_host = ipaddr.resolve_target_addresses(host)
+            family = cls._address_family(args)
+            log.debug("address family: {}".format(family))
+            target_host = ipaddr.resolve_target_addresses(host, family)
             log.debug("target_host for host \"{}\": {}".format(
                 host, target_host))
 
diff --git a/cdist/util/ipaddr.py b/cdist/util/ipaddr.py
index 0bcb0a83..9b730225 100644
--- a/cdist/util/ipaddr.py
+++ b/cdist/util/ipaddr.py
@@ -23,13 +23,13 @@ import socket
 import logging
 
 
-def resolve_target_addresses(host):
-    host_name = resolve_target_host_name(host)
+def resolve_target_addresses(host, family=0):
+    host_name = resolve_target_host_name(host, family)
     host_fqdn = resolve_target_fqdn(host)
     return (host, host_name, host_fqdn)
 
 
-def resolve_target_host_name(host):
+def resolve_target_host_name(host, family=0):
     log = logging.getLogger(host)
     try:
         # getaddrinfo returns a list of 5-tuples:
@@ -38,7 +38,7 @@ def resolve_target_host_name(host):
         # (address, port) for AF_INET,
         # (address, port, flow_info, scopeid) for AF_INET6
         ip_addr = socket.getaddrinfo(
-                host, None, type=socket.SOCK_STREAM)[0][4][0]
+                host, None, family=family, type=socket.SOCK_STREAM)[0][4][0]
         # gethostbyaddr returns triple
         # (hostname, aliaslist, ipaddrlist)
         host_name = socket.gethostbyaddr(ip_addr)[0]
diff --git a/docs/src/man1/cdist.rst b/docs/src/man1/cdist.rst
index 8c6ca549..c76520a6 100644
--- a/docs/src/man1/cdist.rst
+++ b/docs/src/man1/cdist.rst
@@ -20,16 +20,16 @@ SYNOPSIS
                  [-C CACHE_PATH_PATTERN] [-c CONF_DIR] [-i MANIFEST]
                  [-j [JOBS]] [-n] [-o OUT_PATH] [-R [{tar,tgz,tbz2,txz}]]
                  [-r REMOTE_OUT_DIR] [--remote-copy REMOTE_COPY]
-                 [--remote-exec REMOTE_EXEC] [-I INVENTORY_DIR] [-A] [-a]
-                 [-f HOSTFILE] [-p [HOST_MAX]] [-S] [-s] [-t]
+                 [--remote-exec REMOTE_EXEC] [-I INVENTORY_DIR] [-4] [-6]
+                 [-A] [-a] [-f HOSTFILE] [-p [HOST_MAX]] [-S] [-s] [-t]
                  [host [host ...]] 
 
     cdist install [-h] [-l LOGLEVEL] [-q] [-v] [-b] [-g CONFIG_FILE]
                   [-C CACHE_PATH_PATTERN] [-c CONF_DIR] [-i MANIFEST]
                   [-j [JOBS]] [-n] [-o OUT_PATH] [-R [{tar,tgz,tbz2,txz}]]
                   [-r REMOTE_OUT_DIR] [--remote-copy REMOTE_COPY]
-                  [--remote-exec REMOTE_EXEC] [-I INVENTORY_DIR] [-A] [-a]
-                  [-f HOSTFILE] [-p [HOST_MAX]] [-S] [-s] [-t]
+                  [--remote-exec REMOTE_EXEC] [-I INVENTORY_DIR] [-4] [-6]
+                  [-A] [-a] [-f HOSTFILE] [-p [HOST_MAX]] [-S] [-s] [-t]
                   [host [host ...]] 
 
     cdist inventory [-h] [-l LOGLEVEL] [-q] [-v] [-b] [-g CONFIG_FILE]
@@ -118,6 +118,16 @@ CONFIG/INSTALL
 Configure/install one or more hosts.
 Install command is currently in beta.
 
+.. option:: -4, --force-ipv4
+
+    Force to use IPv4 addresses only. No influence for
+    custom remote commands.
+
+.. option:: -6, --force-ipv6
+
+    Force to use IPv6 addresses only. No influence for
+    custom remote commands.
+
 .. option:: -A, --all-tagged
 
     Use all hosts present in tags db. Currently in beta.