Merge pull request #824 from gpotter2/prompt-changer

Apply IPYTHON conf.prompt

scapy/config.py

@@ -14,7 +14,7 @@ import os,time,socket,sys
 from scapy import VERSION
 from scapy.data import *
 from scapy import base_classes
-from scapy.themes import NoTheme, apply_ipython_color
+from scapy.themes import NoTheme, apply_ipython_style
 from scapy.error import log_scapy
 import scapy.modules.six as six
 
@@ -329,11 +329,11 @@ def isCryptographyAdvanced():
 def _prompt_changer(attr, val):
     """Change the current prompt theme"""
     try:
-        sys.ps1 = val.prompt(conf.prompt)
+        sys.ps1 = conf.color_theme.prompt(conf.prompt)
     except:
         pass
     try:
-        apply_ipython_color(get_ipython())
+        apply_ipython_style(get_ipython())
     except NameError:
         pass
 
@@ -385,7 +385,7 @@ debug_tls:When 1, print some TLS session secrets when they are computed.
     checkIPinIP = True
     check_TCPerror_seqack = 0
     verb = 2
-    prompt = ">>> "
+    prompt = Interceptor("prompt", ">>> ", _prompt_changer)
     promisc = 1
     sniff_promisc = 1
     raw_layer = None

scapy/main.py

@@ -26,7 +26,7 @@ except:
 # before the console handlers gets added in interact()
 from scapy.error import log_interactive, log_loading, log_scapy, warning
 import scapy.modules.six as six
-from scapy.themes import DefaultTheme, apply_ipython_color
+from scapy.themes import DefaultTheme, apply_ipython_style
 
 IGNORED = list(six.moves.builtins.__dict__)
 
@@ -200,6 +200,12 @@ def list_contrib(name=None):
 ## Session saving/restoring ##
 ##############################
 
+def update_ipython_session(session):
+    """Updates IPython session with a custom one"""
+    try:
+        get_ipython().user_ns.update(session)
+    except:
+        pass
 
 def save_session(fname=None, session=None, pickleProto=-1):
     """Save current Scapy session to the file specified in the fname arg.
@@ -216,16 +222,21 @@ def save_session(fname=None, session=None, pickleProto=-1):
     log_interactive.info("Use [%s] as session file" % fname)
 
     if session is None:
-        session = six.moves.builtins.__dict__["scapy_session"]
+        try:
+            session = get_ipython().user_ns
+        except:
+            session = six.moves.builtins.__dict__["scapy_session"]
 
     to_be_saved = session.copy()
     if "__builtins__" in to_be_saved:
         del(to_be_saved["__builtins__"])
 
-    for k in to_be_saved.keys():
+    for k in list(to_be_saved.keys()):
         i = to_be_saved[k]
         if hasattr(i, "__module__") and (k[0] == "_" or i.__module__.startswith("IPython")):
             del(to_be_saved[k])
+        if isinstance(i, ConfClass):
+            del(to_be_saved[k])
         elif isinstance(i, (type, type, types.ModuleType)):
             if k[0] != "_":
                 log_interactive.error("[%s] (%s) can't be saved.", k, type(to_be_saved[k]))
@@ -261,6 +272,7 @@ def load_session(fname=None):
     scapy_session = six.moves.builtins.__dict__["scapy_session"]
     scapy_session.clear()
     scapy_session.update(s)
+    update_ipython_session(scapy_session)
 
     log_loading.info("Loaded session [%s]" % fname)
     
@@ -277,6 +289,7 @@ def update_session(fname=None):
         s = six.moves.cPickle.load(open(fname,"rb"))
     scapy_session = six.moves.builtins.__dict__["scapy_session"]
     scapy_session.update(s)
+    update_ipython_session(scapy_session)
 
 def init_session(session_name, mydict=None):
     global SESSION
@@ -319,6 +332,7 @@ def init_session(session_name, mydict=None):
 
     if mydict is not None:
         six.moves.builtins.__dict__["scapy_session"].update(mydict)
+        update_ipython_session(mydict)
         GLOBKEYS.extend(mydict)
 
 ################
@@ -489,13 +503,7 @@ def interact(mydict=None,argv=None,mybanner=None,loglevel=20):
             get_ipython
         except NameError:
             # Set "classic" prompt style when launched from run_scapy(.bat) files
-            class ClassicPrompt(Prompts):
-                def in_prompt_tokens(self, cli=None):
-                   return [(Token.Prompt, '>>> '),]
-                def out_prompt_tokens(self):
-                   return [(Token.OutPrompt, ''),]
-            cfg.TerminalInteractiveShell.prompts_class=ClassicPrompt # Set classic prompt style
-            apply_ipython_color(shell=cfg.TerminalInteractiveShell) # Register and apply scapy color style
+            apply_ipython_style(shell=cfg.TerminalInteractiveShell) # Register and apply scapy color+prompt style
             cfg.TerminalInteractiveShell.confirm_exit = False # Remove confirm exit
             cfg.TerminalInteractiveShell.separate_in = u'' # Remove spacing line
 

scapy/themes.py

@@ -288,11 +288,21 @@ class HTMLTheme2(HTMLTheme):
     style_right = "#[#span class=right#]#%s#[#/span#]#"
 
 
-def apply_ipython_color(shell):
+def apply_ipython_style(shell):
     """Updates the specified IPython console shell with
     the conf.color_theme scapy theme."""
     from IPython.terminal.prompts import Prompts, Token
-    shell.highlighting_style_overrides = {
+    from scapy.config import conf
+    if isinstance(conf.prompt, Prompts):
+        shell.prompts_class = conf.prompt # Set custom prompt style
+    else:
+        class ClassicPrompt(Prompts):
+            def in_prompt_tokens(self, cli=None):
+               return [(Token.Prompt, str(conf.prompt)),]
+            def out_prompt_tokens(self):
+               return [(Token.OutPrompt, ''),]
+        shell.prompts_class=ClassicPrompt # Apply classic prompt style
+    shell.highlighting_style_overrides = { # Register and apply scapy color style
         Token.Prompt: Color.ansi_to_pygments(conf.color_theme.style_prompt),
     }
     try: