From 9408cbc5145022833cb4f526508735f18a58108d Mon Sep 17 00:00:00 2001
From: Bilal Amarni <bilal.amarni@gmail.com>
Date: Thu, 5 Apr 2018 14:27:13 +0200
Subject: [PATCH] Rewrite Cosmos binary resource URLs when needed

Depending on the user setup, they might contain an incorrect scheme.
This is because Cosmos relies in the scheme used to query admin router.
If a CLI is configured to use HTTPS but a load balancer sits in front
of admin router and does TLS termination, these URLs will use the HTTP
scheme.

Not only that it is a security concern, it would simply fail on setups
which don't support HTTPS. This is hopefully a temporary fix until we
come-up with a long-term fix in DC/OS.

Fixes #1190
https://jira.mesosphere.com/browse/COPS-3052
https://jira.mesosphere.com/browse/DCOS_OSS-2325
---
 dcos/subcommand.py       | 31 +++++++++++++++++++++++++++++++
 tests/test_subcommand.py | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/dcos/subcommand.py b/dcos/subcommand.py
index b45d4a47a..89ed1ebff 100644
--- a/dcos/subcommand.py
+++ b/dcos/subcommand.py
@@ -9,7 +9,9 @@
 import subprocess
 import sys
 import zipfile
+
 from distutils.version import LooseVersion
+from urllib.parse import urlparse
 
 from dcos import config, constants, http, util
 from dcos.errors import DCOSException
@@ -515,6 +517,10 @@ def _install_with_binary(
 
     binary_url, kind = binary_cli.get("url"), binary_cli.get("kind")
 
+    binary_url = _rewrite_binary_url(
+        binary_url,
+        config.get_config_val("core.dcos_url"))
+
     try:
         env_bin_dir = os.path.join(env_directory, BIN_DIRECTORY)
 
@@ -567,6 +573,31 @@ def _install_with_binary(
     return None
 
 
+def _rewrite_binary_url(binary_url, dcos_url):
+    """ Rewrite Cosmos binary URLs when needed.
+
+    A rewrite happens when a resource refers to the currently configured
+    DC/OS cluster but uses the wrong scheme. Ideally this "hack" will be
+    removed in the future.
+
+    For more context, see https://jira.mesosphere.com/browse/DCOS_OSS-2325.
+
+    :param binary_url: the binary resource URL as returned by Cosmos
+    :type binary_url: str
+    :param dcos_url: the DC/OS URL of the current cluster
+    :type dcos_url: str
+    :rtype: str
+    """
+
+    binary_url = urlparse(binary_url)
+    dcos_url = urlparse(dcos_url)
+
+    if binary_url.netloc == dcos_url.netloc:
+        binary_url = binary_url._replace(scheme=dcos_url.scheme)
+
+    return binary_url.geturl()
+
+
 def _install_with_pip(
         package_name,
         env_directory,
diff --git a/tests/test_subcommand.py b/tests/test_subcommand.py
index 0eec17326..2af18d3cd 100644
--- a/tests/test_subcommand.py
+++ b/tests/test_subcommand.py
@@ -7,3 +7,37 @@ def test_noun():
 
 def test_hyphen_noun():
     assert subcommand.noun("some/path/to/dcos-sub-command") == "sub-command"
+
+
+def test_rewrite_binary_url_with_incorrect_scheme():
+    binary_url = ("http://dcos.example.com/package/resource"
+                  "?url=https://binary.example.com")
+
+    dcos_url = "https://dcos.example.com"
+
+    rewritten_url = subcommand._rewrite_binary_url(binary_url, dcos_url)
+
+    assert rewritten_url == ("https://dcos.example.com/package/resource"
+                             "?url=https://binary.example.com")
+
+
+def test_rewrite_binary_url_with_correct_scheme():
+    binary_url = ("https://dcos.example.com/package/resource"
+                  "?url=https://binary.example.com")
+
+    dcos_url = "https://dcos.example.com"
+
+    rewritten_url = subcommand._rewrite_binary_url(binary_url, dcos_url)
+
+    assert rewritten_url == binary_url
+
+
+def test_rewrite_binary_url_with_external_url():
+    binary_url = ("https://not-dcos.example.com/package/resource"
+                  "?url=https://binary.example.com")
+
+    dcos_url = "https://dcos.example.com"
+
+    rewritten_url = subcommand._rewrite_binary_url(binary_url, dcos_url)
+
+    assert rewritten_url == binary_url