From b362adf7ca5c494aecdbce827ef6e3044194f38e Mon Sep 17 00:00:00 2001
From: Jaye Doepke <jdoepke@mintel.com>
Date: Mon, 2 Oct 2017 17:24:08 -0500
Subject: [PATCH] Relay and transit url conf from env vars

Added the click option to look for relay and transit urls in environment
variables. If you're running your own relay/transit servers (such as
inside a corporate firewall), this will make client's lives easier.
---
 src/wormhole/cli/cli.py        |  2 ++
 src/wormhole/test/test_args.py | 51 ++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/src/wormhole/cli/cli.py b/src/wormhole/cli/cli.py
index 38851d2..58ccd8f 100644
--- a/src/wormhole/cli/cli.py
+++ b/src/wormhole/cli/cli.py
@@ -60,11 +60,13 @@ class AliasedGroup(click.Group):
     "--appid", default=None, metavar="APPID", help="appid to use")
 @click.option(
     "--relay-url", default=public_relay.RENDEZVOUS_RELAY,
+    envvar='WORMHOLE_RELAY_URL',
     metavar="URL",
     help="rendezvous relay to use",
 )
 @click.option(
     "--transit-helper", default=public_relay.TRANSIT_RELAY,
+    envvar='WORMHOLE_TRANSIT_HELPER',
     metavar="tcp:HOST:PORT",
     help="transit relay to use",
 )
diff --git a/src/wormhole/test/test_args.py b/src/wormhole/test/test_args.py
index c013e29..ff6f060 100644
--- a/src/wormhole/test/test_args.py
+++ b/src/wormhole/test/test_args.py
@@ -1,4 +1,6 @@
+import os
 import sys
+import mock
 from twisted.trial import unittest
 from ..cli.public_relay import RENDEZVOUS_RELAY, TRANSIT_RELAY
 from .common import config
@@ -71,6 +73,31 @@ class Send(unittest.TestCase):
         cfg = config("send", "-0", "fn")
         self.assertEqual(cfg.zeromode, True)
 
+    def test_relay_env_var(self):
+        relay_url = str(mock.sentinel.relay_url)
+        with mock.patch.dict(os.environ, WORMHOLE_RELAY_URL=relay_url):
+            cfg = config("send")
+        self.assertEqual(cfg.relay_url, relay_url)
+
+        # Make sure cmd line option overrides environment variable
+        relay_url_2 = str(mock.sentinel.relay_url_2)
+        with mock.patch.dict(os.environ, WORMHOLE_RELAY_URL=relay_url):
+            cfg = config("--relay-url", relay_url_2, "send")
+        self.assertEqual(cfg.relay_url, relay_url_2)
+
+    def test_transit_env_var(self):
+        transit_url = str(mock.sentinel.transit_url)
+        with mock.patch.dict(os.environ, WORMHOLE_TRANSIT_HELPER=transit_url):
+            cfg = config("send")
+        self.assertEqual(cfg.transit_helper, transit_url)
+
+        # Make sure cmd line option overrides environment variable
+        transit_url_2 = str(mock.sentinel.transit_url_2)
+        with mock.patch.dict(os.environ, WORMHOLE_TRANSIT_HELPER=transit_url):
+            cfg = config("--transit-helper", transit_url_2, "send")
+        self.assertEqual(cfg.transit_helper, transit_url_2)
+
+
 class Receive(unittest.TestCase):
     def test_baseline(self):
         cfg = config("receive")
@@ -140,6 +167,30 @@ class Receive(unittest.TestCase):
         cfg = config("receive", "--output-file", "fn")
         self.assertEqual(cfg.output_file, u"fn")
 
+    def test_relay_env_var(self):
+        relay_url = str(mock.sentinel.relay_url)
+        with mock.patch.dict(os.environ, WORMHOLE_RELAY_URL=relay_url):
+            cfg = config("receive")
+        self.assertEqual(cfg.relay_url, relay_url)
+
+        # Make sure cmd line option overrides environment variable
+        relay_url_2 = str(mock.sentinel.relay_url_2)
+        with mock.patch.dict(os.environ, WORMHOLE_RELAY_URL=relay_url):
+            cfg = config("--relay-url", relay_url_2, "receive")
+        self.assertEqual(cfg.relay_url, relay_url_2)
+
+    def test_transit_env_var(self):
+        transit_url = str(mock.sentinel.transit_url)
+        with mock.patch.dict(os.environ, WORMHOLE_TRANSIT_HELPER=transit_url):
+            cfg = config("receive")
+        self.assertEqual(cfg.transit_helper, transit_url)
+
+        # Make sure cmd line option overrides environment variable
+        transit_url_2 = str(mock.sentinel.transit_url_2)
+        with mock.patch.dict(os.environ, WORMHOLE_TRANSIT_HELPER=transit_url):
+            cfg = config("--transit-helper", transit_url_2, "receive")
+        self.assertEqual(cfg.transit_helper, transit_url_2)
+
 class Config(unittest.TestCase):
     def test_send(self):
         cfg = config("send")