diff --git a/scapy/config.py b/scapy/config.py index cc97c6dcf..4e482c400 100755 --- a/scapy/config.py +++ b/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 diff --git a/scapy/main.py b/scapy/main.py index af9bb1b4a..d4b30c39e 100644 --- a/scapy/main.py +++ b/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 diff --git a/scapy/themes.py b/scapy/themes.py index f613c6953..173194ef4 100644 --- a/scapy/themes.py +++ b/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: