From 25eb000cf8412213fde768eafd6a0d573fd992cc Mon Sep 17 00:00:00 2001
From: Chad Wilson <chadw@thoughtworks.com>
Date: Mon, 26 Oct 2020 01:12:53 +0800
Subject: [PATCH] Support `HTTPS_PROXY` (or `HTTP_PROXY`) env vars for
 specifying a proxy to route traffic through.

Note that this will only work with a TLS passthrough proxy anyway; since the TLS connections will fail cert validation otherwise.
---
 eksrollup/lib/k8s.py                    |  6 ++++++
 requirements-tests.txt                  |  2 --
 tests/fixtures/example-kube-config.yaml | 16 ++++++++++++++++
 tests/test_k8s.py                       | 21 +++++++++++++++++++--
 4 files changed, 41 insertions(+), 4 deletions(-)
 create mode 100644 tests/fixtures/example-kube-config.yaml

diff --git a/eksrollup/lib/k8s.py b/eksrollup/lib/k8s.py
index 37c650f..11d5d0f 100644
--- a/eksrollup/lib/k8s.py
+++ b/eksrollup/lib/k8s.py
@@ -1,5 +1,6 @@
 from kubernetes import client, config
 from kubernetes.client.rest import ApiException
+import os
 import subprocess
 import time
 import sys
@@ -16,6 +17,11 @@ def ensure_config_loaded():
         except config.ConfigException:
             raise Exception("Could not configure kubernetes python client")
 
+    proxy_url = os.getenv('HTTPS_PROXY', os.getenv('HTTP_PROXY', None))
+    if proxy_url:
+        logger.info(f"Setting proxy: {proxy_url}")
+        client.Configuration._default.proxy = proxy_url
+
 
 def get_k8s_nodes(exclude_node_label_keys=app_config["EXCLUDE_NODE_LABEL_KEYS"]):
     """
diff --git a/requirements-tests.txt b/requirements-tests.txt
index 959ecd6..e24d17c 100644
--- a/requirements-tests.txt
+++ b/requirements-tests.txt
@@ -2,6 +2,4 @@ moto==1.3.13
 python-box~=3.4.0
 flake8~=3.7.7
 nose2~=0.9.1
-kubernetes~=10.0.1
-python-dotenv~=0.10.2
 twine
\ No newline at end of file
diff --git a/tests/fixtures/example-kube-config.yaml b/tests/fixtures/example-kube-config.yaml
new file mode 100644
index 0000000..b3e2b59
--- /dev/null
+++ b/tests/fixtures/example-kube-config.yaml
@@ -0,0 +1,16 @@
+apiVersion: v1
+kind: Config
+current-context: example-context
+clusters:
+  - cluster:
+      server: https://localhost:8080
+    name: local-cluster
+contexts:
+  - context:
+      cluster: local-cluster
+      namespace: default
+      user: user
+    name: example-context
+users:
+  - name: user
+    user: {}
\ No newline at end of file
diff --git a/tests/test_k8s.py b/tests/test_k8s.py
index e1cabd1..3c8d443 100644
--- a/tests/test_k8s.py
+++ b/tests/test_k8s.py
@@ -1,10 +1,11 @@
 import os
 import unittest
 import json
-from eksrollup.lib.k8s import k8s_nodes_count, k8s_nodes_ready, get_node_by_instance_id
+from eksrollup.lib.k8s import k8s_nodes_count, k8s_nodes_ready, get_node_by_instance_id, ensure_config_loaded
 from unittest.mock import patch
 from box import Box
-
+from kubernetes.client import ApiClient
+from kubernetes.config import kube_config
 
 class TestK8S(unittest.TestCase):
 
@@ -17,6 +18,22 @@ def setUp(self):
         with open(f"{current_dir}/fixtures/k8s_response_unhealthy.json", "r") as file:
             self.k8s_response_mock_unhealthy = json.load(file)
 
+        kube_config.KUBE_CONFIG_DEFAULT_LOCATION = f"{current_dir}/fixtures/example-kube-config.yaml"
+
+    def test_ensure_config_loaded_proxy_default(self):
+        ensure_config_loaded()
+        self.assertEqual(None, ApiClient().configuration.proxy)
+
+    def test_ensure_config_loaded_proxy_prefers_https(self):
+        with patch.dict(os.environ, {'HTTPS_PROXY': 'http://localhost:12345', 'HTTP_PROXY': 'http://localhost:6789'}):
+            ensure_config_loaded()
+            self.assertEqual('http://localhost:12345', ApiClient().configuration.proxy)
+
+    def test_ensure_config_loaded_proxy_fallback_to_http(self):
+        with patch.dict(os.environ, {'HTTP_PROXY': 'http://localhost:12345'}):
+            ensure_config_loaded()
+            self.assertEqual('http://localhost:12345', ApiClient().configuration.proxy)
+
     def test_k8s_node_count(self):
         with patch('eksrollup.lib.k8s.get_k8s_nodes') as get_k8s_nodes_mock:
             get_k8s_nodes_mock.return_value = self.k8s_response_mock['items']