From 18fe02f7e7abe3d399a8b8f89c53fd2334ce4c25 Mon Sep 17 00:00:00 2001
From: gpotter2 <gpotter2@users.noreply.github.com>
Date: Sun, 15 Jan 2017 18:32:00 +0100
Subject: [PATCH] [Windows] Fix AutoCompletion + add tests (#465)

---
 appveyor.yml                   |  2 +-
 scapy/arch/windows/__init__.py | 17 ++++++++
 scapy/config.py                |  1 +
 scapy/main.py                  |  2 +-
 test/mock_windows.uts          | 78 ++++++++++++++++++++++++++++++++++
 5 files changed, 98 insertions(+), 2 deletions(-)

diff --git a/appveyor.yml b/appveyor.yml
index 1686a008..fe8d846e 100644
--- a/appveyor.yml
+++ b/appveyor.yml
@@ -19,7 +19,7 @@ install:
   - refreshenv
 
   # Install Python modules
-  - "%PYTHON%\\python -m pip install ecdsa cryptography coverage mock"
+  - "%PYTHON%\\python -m pip install ecdsa cryptography coverage mock pyreadline keyboard"
 
 test_script:
   # Set environment variables
diff --git a/scapy/arch/windows/__init__.py b/scapy/arch/windows/__init__.py
index e60998ba..e92fdcf6 100755
--- a/scapy/arch/windows/__init__.py
+++ b/scapy/arch/windows/__init__.py
@@ -550,12 +550,29 @@ if conf.interactive_shell != 'ipython':
     try:
         __IPYTHON__
     except NameError:
+        def readLineScapy(prompt):
+            result = ""
+            end = False
+            while not end :
+                if not end and result != "":
+                    line = readline.rl.readline("... ")
+                else:
+                    line = readline.rl.readline(prompt)
+                if line.strip().endswith(":"):
+                    end = False
+                elif result == "":
+                    end = True
+                if line.strip() == "":
+                    end = True
+                result = result + "\n" + line
+            return unicode(result)
         try:
             import readline
             console = readline.GetOutputFile()
         except (ImportError, AttributeError):
             log_loading.info("Could not get readline console. Will not interpret ANSI color codes.") 
         else:
+            conf.readfunc = readLineScapy
             orig_stdout = sys.stdout
             sys.stdout = console
 
diff --git a/scapy/config.py b/scapy/config.py
index 5b7419aa..b97ee0bd 100755
--- a/scapy/config.py
+++ b/scapy/config.py
@@ -333,6 +333,7 @@ contribs: a dict which can be used by contrib layers to store local configuratio
     stealth = "not implemented"
     iface = None
     iface6 = None
+    readfunc = None
     layers = LayersList()
     commands = CommandsList()
     logLevel = LogLevel()
diff --git a/scapy/main.py b/scapy/main.py
index edce92be..e788a2fc 100644
--- a/scapy/main.py
+++ b/scapy/main.py
@@ -381,7 +381,7 @@ def interact(mydict=None,argv=None,mybanner=None,loglevel=20):
 
     else:
         code.interact(banner = the_banner % (conf.version),
-                      local=session)
+                      local=session, readfunc=conf.readfunc)
 
     if conf.session:
         save_session(conf.session, session)
diff --git a/test/mock_windows.uts b/test/mock_windows.uts
index 25d495cb..191938d0 100644
--- a/test/mock_windows.uts
+++ b/test/mock_windows.uts
@@ -69,3 +69,81 @@ ifIndex DestinationPrefix                          NextHop
 
 
 test_read_routes6_windows()
+
+############
+############
++ Main.py emulator
+
+= Prepare readline patching functions
+
+from scapy.main import *
+import scapy.config as conf
+import sys
+
+import mock
+import readline
+from threading import Thread, Event
+
+class sendTextAndTab(Thread):
+    """Send text directly as Input"""
+    def __init__(self, event, text):
+        Thread.__init__(self)
+        self.stopped = event
+        self.send_text = text
+    def run(self):
+        import keyboard
+        time.sleep(1)
+        while not self.stopped.wait(0.5):
+            keyboard.write(self.send_text)
+            keyboard.send("tab")
+            keyboard.send("enter")
+
+index = 0
+@mock.patch("pyreadline.console.console.Console.size")
+@mock.patch("scapy.config.conf.readfunc")
+def emulate_main_input(data, mock_readfunc, mock_pyr_size):
+    # This fix when the windows doesn't have a size (run with a windows server)
+    mock_pyr_size.return_value = (300, 300)
+    global index
+    index = 0 # reset var
+    def readlineScapy(*args, **kargs):
+        global index
+        if len(data) == index:
+            r_data = "exit(1 if hasattr(sys, 'last_value') and sys.last_value is not None else 0)"
+        else:
+            r_data = data[index]
+            if r_data.startswith("#AUTOCOMPLETE"):
+                send_text = re.match(r'#AUTOCOMPLETE{(.*)}', r_data).group(1)
+                stopFlag = Event()
+                thread = sendTextAndTab(stopFlag, send_text)
+                thread.start()
+                # This will block the program until the thread has pushed the stuff
+                r_data = readline.rl.readline()
+                stopFlag.set()
+        index +=1
+        print r_data
+        return r_data
+    mock_readfunc.side_effect = readlineScapy
+    sys.argv = ['']
+    def console_exit(code):
+        raise SystemExit(code)
+    exit_code = -1
+    try:
+        interact(mydict={"exit": console_exit})
+    except SystemExit as e:
+        exit_code = str(e)
+        pass
+    assert exit_code == "0"
+
+= Test basic running
+data = ["IP()", "assert _.name == 'IP'"]
+emulate_main_input(data)
+
+= Test function parsing
+data = ["def test():", "    return True", "", "assert test() == True"]
+emulate_main_input(data)
+
+= Test auto-completion
+from ctypes import wintypes
+data = ["#AUTOCOMPLETE{scapy.config.conf.vers}", "assert _ == scapy.config.conf.version"]
+emulate_main_input(data)
-- 
GitLab