• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# SPDX-License-Identifier: GPL-2.0-only
2# This file is part of Scapy
3# See https://scapy.net/ for more information
4# Copyright (C) Philippe Biondi <phil@secdev.org>
5
6"""
7Color themes for the interactive console.
8"""
9
10##################
11#  Color themes  #
12##################
13
14import html
15import sys
16
17from typing import (
18    Any,
19    List,
20    Optional,
21    Tuple,
22    cast,
23)
24from scapy.compat import Protocol
25
26
27class ColorTable:
28    colors = {  # Format: (ansi, pygments)
29        # foreground
30        "black": ("\033[30m", "#ansiblack"),
31        "red": ("\033[31m", "#ansired"),
32        "green": ("\033[32m", "#ansigreen"),
33        "yellow": ("\033[33m", "#ansiyellow"),
34        "blue": ("\033[34m", "#ansiblue"),
35        "purple": ("\033[35m", "#ansipurple"),
36        "cyan": ("\033[36m", "#ansicyan"),
37        "white": ("\033[37m", "#ansiwhite"),
38        "grey": ("\033[38;5;246m", "#ansiwhite"),
39        "reset": ("\033[39m", "noinherit"),
40        # background
41        "bg_black": ("\033[40m", "bg:#ansiblack"),
42        "bg_red": ("\033[41m", "bg:#ansired"),
43        "bg_green": ("\033[42m", "bg:#ansigreen"),
44        "bg_yellow": ("\033[43m", "bg:#ansiyellow"),
45        "bg_blue": ("\033[44m", "bg:#ansiblue"),
46        "bg_purple": ("\033[45m", "bg:#ansipurple"),
47        "bg_cyan": ("\033[46m", "bg:#ansicyan"),
48        "bg_white": ("\033[47m", "bg:#ansiwhite"),
49        "bg_reset": ("\033[49m", "noinherit"),
50        # specials
51        "normal": ("\033[0m", "noinherit"),  # color & brightness
52        "bold": ("\033[1m", "bold"),
53        "uline": ("\033[4m", "underline"),
54        "blink": ("\033[5m", ""),
55        "invert": ("\033[7m", ""),
56    }
57    inv_map = {v[0]: v[1] for k, v in colors.items()}
58
59    def __repr__(self):
60        # type: () -> str
61        return "<ColorTable>"
62
63    def __getattr__(self, attr):
64        # type: (str) -> str
65        return self.colors.get(attr, [""])[0]
66
67    def ansi_to_pygments(self, x):
68        # type: (str) -> str
69        """
70        Transform ansi encoded text to Pygments text
71        """
72        for k, v in self.inv_map.items():
73            x = x.replace(k, " " + v)
74        return x.strip()
75
76
77Color = ColorTable()
78
79
80class _ColorFormatterType(Protocol):
81    def __call__(self,
82                 val: Any,
83                 fmt: Optional[str] = None,
84                 fmt2: str = "",
85                 before: str = "",
86                 after: str = "") -> str:
87        pass
88
89
90def create_styler(fmt=None,  # type: Optional[str]
91                  before="",  # type: str
92                  after="",  # type: str
93                  fmt2="%s"  # type: str
94                  ):
95    # type: (...) -> _ColorFormatterType
96    def do_style(val: Any,
97                 fmt: Optional[str] = fmt,
98                 fmt2: str = fmt2,
99                 before: str = before,
100                 after: str = after) -> str:
101        if fmt is None:
102            sval = str(val)
103        else:
104            sval = fmt % val
105        return fmt2 % (before + sval + after)
106    return do_style
107
108
109class ColorTheme:
110    style_normal = ""
111    style_prompt = ""
112    style_punct = ""
113    style_id = ""
114    style_not_printable = ""
115    style_layer_name = ""
116    style_field_name = ""
117    style_field_value = ""
118    style_emph_field_name = ""
119    style_emph_field_value = ""
120    style_depreciate_field_name = ""
121    style_packetlist_name = ""
122    style_packetlist_proto = ""
123    style_packetlist_value = ""
124    style_fail = ""
125    style_success = ""
126    style_odd = ""
127    style_even = ""
128    style_opening = ""
129    style_active = ""
130    style_closed = ""
131    style_left = ""
132    style_right = ""
133    style_logo = ""
134
135    def __repr__(self):
136        # type: () -> str
137        return "<%s>" % self.__class__.__name__
138
139    def __reduce__(self):
140        # type: () -> Tuple[type, Any, Any]
141        return (self.__class__, (), ())
142
143    def __getattr__(self, attr):
144        # type: (str) -> _ColorFormatterType
145        if attr in ["__getstate__", "__setstate__", "__getinitargs__",
146                    "__reduce_ex__"]:
147            raise AttributeError()
148        return create_styler()
149
150    def format(self, string, fmt):
151        # type: (str, str) -> str
152        for style in fmt.split("+"):
153            string = getattr(self, style)(string)
154        return string
155
156
157class NoTheme(ColorTheme):
158    pass
159
160
161class AnsiColorTheme(ColorTheme):
162    def __getattr__(self, attr):
163        # type: (str) -> _ColorFormatterType
164        if attr.startswith("__"):
165            raise AttributeError(attr)
166        s = "style_%s" % attr
167        if s in self.__class__.__dict__:
168            before = getattr(self, s)
169            after = self.style_normal
170        elif not isinstance(self, BlackAndWhite) and attr in Color.colors:
171            before = Color.colors[attr][0]
172            after = Color.colors["normal"][0]
173        else:
174            before = after = ""
175
176        return create_styler(before=before, after=after)
177
178
179class BlackAndWhite(AnsiColorTheme, NoTheme):
180    pass
181
182
183class DefaultTheme(AnsiColorTheme):
184    style_normal = Color.normal
185    style_prompt = Color.blue + Color.bold
186    style_punct = Color.normal
187    style_id = Color.blue + Color.bold
188    style_not_printable = Color.white
189    style_depreciate_field_name = Color.grey
190    style_layer_name = Color.red + Color.bold
191    style_field_name = Color.blue
192    style_field_value = Color.purple
193    style_emph_field_name = Color.blue + Color.uline + Color.bold
194    style_emph_field_value = Color.purple + Color.uline + Color.bold
195    style_packetlist_name = Color.red + Color.bold
196    style_packetlist_proto = Color.blue
197    style_packetlist_value = Color.purple
198    style_fail = Color.red + Color.bold
199    style_success = Color.blue + Color.bold
200    style_even = Color.black + Color.bold
201    style_odd = Color.black
202    style_opening = Color.yellow
203    style_active = Color.black
204    style_closed = Color.white
205    style_left = Color.blue + Color.invert
206    style_right = Color.red + Color.invert
207    style_logo = Color.green + Color.bold
208
209
210class BrightTheme(AnsiColorTheme):
211    style_normal = Color.normal
212    style_punct = Color.normal
213    style_id = Color.yellow + Color.bold
214    style_layer_name = Color.red + Color.bold
215    style_field_name = Color.yellow + Color.bold
216    style_field_value = Color.purple + Color.bold
217    style_emph_field_name = Color.yellow + Color.bold
218    style_emph_field_value = Color.green + Color.bold
219    style_packetlist_name = Color.red + Color.bold
220    style_packetlist_proto = Color.yellow + Color.bold
221    style_packetlist_value = Color.purple + Color.bold
222    style_fail = Color.red + Color.bold
223    style_success = Color.blue + Color.bold
224    style_even = Color.black + Color.bold
225    style_odd = Color.black
226    style_left = Color.cyan + Color.invert
227    style_right = Color.purple + Color.invert
228    style_logo = Color.green + Color.bold
229
230
231class RastaTheme(AnsiColorTheme):
232    style_normal = Color.normal + Color.green + Color.bold
233    style_prompt = Color.yellow + Color.bold
234    style_punct = Color.red
235    style_id = Color.green + Color.bold
236    style_not_printable = Color.green
237    style_layer_name = Color.red + Color.bold
238    style_field_name = Color.yellow + Color.bold
239    style_field_value = Color.green + Color.bold
240    style_emph_field_name = Color.green
241    style_emph_field_value = Color.green
242    style_packetlist_name = Color.red + Color.bold
243    style_packetlist_proto = Color.yellow + Color.bold
244    style_packetlist_value = Color.green + Color.bold
245    style_fail = Color.red
246    style_success = Color.red + Color.bold
247    style_even = Color.yellow
248    style_odd = Color.green
249    style_left = Color.yellow + Color.invert
250    style_right = Color.red + Color.invert
251    style_logo = Color.green + Color.bold
252
253
254class ColorOnBlackTheme(AnsiColorTheme):
255    """Color theme for black backgrounds"""
256    style_normal = Color.normal
257    style_prompt = Color.green + Color.bold
258    style_punct = Color.normal
259    style_id = Color.green
260    style_not_printable = Color.black + Color.bold
261    style_layer_name = Color.yellow + Color.bold
262    style_field_name = Color.cyan
263    style_field_value = Color.purple + Color.bold
264    style_emph_field_name = Color.cyan + Color.bold
265    style_emph_field_value = Color.red + Color.bold
266    style_packetlist_name = Color.black + Color.bold
267    style_packetlist_proto = Color.yellow + Color.bold
268    style_packetlist_value = Color.purple + Color.bold
269    style_fail = Color.red + Color.bold
270    style_success = Color.green
271    style_even = Color.black + Color.bold
272    style_odd = Color.white
273    style_opening = Color.yellow
274    style_active = Color.white + Color.bold
275    style_closed = Color.black + Color.bold
276    style_left = Color.cyan + Color.bold
277    style_right = Color.red + Color.bold
278    style_logo = Color.green + Color.bold
279
280
281class FormatTheme(ColorTheme):
282    def __getattr__(self, attr: str) -> _ColorFormatterType:
283        if attr.startswith("__"):
284            raise AttributeError(attr)
285        colfmt = self.__class__.__dict__.get("style_%s" % attr, "%s")
286        return create_styler(fmt2=colfmt)
287
288
289class LatexTheme(FormatTheme):
290    r"""
291    You can prepend the output from this theme with
292    \tt\obeyspaces\obeylines\tiny\noindent
293    """
294    style_prompt = r"\textcolor{blue}{%s}"
295    style_not_printable = r"\textcolor{gray}{%s}"
296    style_layer_name = r"\textcolor{red}{\bf %s}"
297    style_field_name = r"\textcolor{blue}{%s}"
298    style_field_value = r"\textcolor{purple}{%s}"
299    style_emph_field_name = r"\textcolor{blue}{\underline{%s}}"  # ul
300    style_emph_field_value = r"\textcolor{purple}{\underline{%s}}"  # ul
301    style_packetlist_name = r"\textcolor{red}{\bf %s}"
302    style_packetlist_proto = r"\textcolor{blue}{%s}"
303    style_packetlist_value = r"\textcolor{purple}{%s}"
304    style_fail = r"\textcolor{red}{\bf %s}"
305    style_success = r"\textcolor{blue}{\bf %s}"
306    style_left = r"\textcolor{blue}{%s}"
307    style_right = r"\textcolor{red}{%s}"
308#    style_even = r"}{\bf "
309#    style_odd = ""
310    style_logo = r"\textcolor{green}{\bf %s}"
311
312    def __getattr__(self, attr: str) -> _ColorFormatterType:
313        from scapy.utils import tex_escape
314        styler = super(LatexTheme, self).__getattr__(attr)
315        return cast(
316            _ColorFormatterType,
317            lambda x, *args, **kwargs: styler(tex_escape(str(x)), *args, **kwargs),
318        )
319
320
321class LatexTheme2(FormatTheme):
322    style_prompt = r"@`@textcolor@[@blue@]@@[@%s@]@"
323    style_not_printable = r"@`@textcolor@[@gray@]@@[@%s@]@"
324    style_layer_name = r"@`@textcolor@[@red@]@@[@@`@bfseries@[@@]@%s@]@"
325    style_field_name = r"@`@textcolor@[@blue@]@@[@%s@]@"
326    style_field_value = r"@`@textcolor@[@purple@]@@[@%s@]@"
327    style_emph_field_name = r"@`@textcolor@[@blue@]@@[@@`@underline@[@%s@]@@]@"
328    style_emph_field_value = r"@`@textcolor@[@purple@]@@[@@`@underline@[@%s@]@@]@"  # noqa: E501
329    style_packetlist_name = r"@`@textcolor@[@red@]@@[@@`@bfseries@[@@]@%s@]@"
330    style_packetlist_proto = r"@`@textcolor@[@blue@]@@[@%s@]@"
331    style_packetlist_value = r"@`@textcolor@[@purple@]@@[@%s@]@"
332    style_fail = r"@`@textcolor@[@red@]@@[@@`@bfseries@[@@]@%s@]@"
333    style_success = r"@`@textcolor@[@blue@]@@[@@`@bfseries@[@@]@%s@]@"
334    style_even = r"@`@textcolor@[@gray@]@@[@@`@bfseries@[@@]@%s@]@"
335#    style_odd = r"@`@textcolor@[@black@]@@[@@`@bfseries@[@@]@%s@]@"
336    style_left = r"@`@textcolor@[@blue@]@@[@%s@]@"
337    style_right = r"@`@textcolor@[@red@]@@[@%s@]@"
338    style_logo = r"@`@textcolor@[@green@]@@[@@`@bfseries@[@@]@%s@]@"
339
340
341class HTMLTheme(FormatTheme):
342    style_prompt = "<span class=prompt>%s</span>"
343    style_not_printable = "<span class=not_printable>%s</span>"
344    style_layer_name = "<span class=layer_name>%s</span>"
345    style_field_name = "<span class=field_name>%s</span>"
346    style_field_value = "<span class=field_value>%s</span>"
347    style_emph_field_name = "<span class=emph_field_name>%s</span>"
348    style_emph_field_value = "<span class=emph_field_value>%s</span>"
349    style_packetlist_name = "<span class=packetlist_name>%s</span>"
350    style_packetlist_proto = "<span class=packetlist_proto>%s</span>"
351    style_packetlist_value = "<span class=packetlist_value>%s</span>"
352    style_fail = "<span class=fail>%s</span>"
353    style_success = "<span class=success>%s</span>"
354    style_even = "<span class=even>%s</span>"
355    style_odd = "<span class=odd>%s</span>"
356    style_left = "<span class=left>%s</span>"
357    style_right = "<span class=right>%s</span>"
358
359
360class HTMLTheme2(HTMLTheme):
361    style_prompt = "#[#span class=prompt#]#%s#[#/span#]#"
362    style_not_printable = "#[#span class=not_printable#]#%s#[#/span#]#"
363    style_layer_name = "#[#span class=layer_name#]#%s#[#/span#]#"
364    style_field_name = "#[#span class=field_name#]#%s#[#/span#]#"
365    style_field_value = "#[#span class=field_value#]#%s#[#/span#]#"
366    style_emph_field_name = "#[#span class=emph_field_name#]#%s#[#/span#]#"
367    style_emph_field_value = "#[#span class=emph_field_value#]#%s#[#/span#]#"
368    style_packetlist_name = "#[#span class=packetlist_name#]#%s#[#/span#]#"
369    style_packetlist_proto = "#[#span class=packetlist_proto#]#%s#[#/span#]#"
370    style_packetlist_value = "#[#span class=packetlist_value#]#%s#[#/span#]#"
371    style_fail = "#[#span class=fail#]#%s#[#/span#]#"
372    style_success = "#[#span class=success#]#%s#[#/span#]#"
373    style_even = "#[#span class=even#]#%s#[#/span#]#"
374    style_odd = "#[#span class=odd#]#%s#[#/span#]#"
375    style_left = "#[#span class=left#]#%s#[#/span#]#"
376    style_right = "#[#span class=right#]#%s#[#/span#]#"
377
378
379def apply_ipython_style(shell):
380    # type: (Any) -> None
381    """Updates the specified IPython console shell with
382    the conf.color_theme scapy theme."""
383    try:
384        from IPython.terminal.prompts import Prompts, Token
385    except Exception:
386        from scapy.error import log_loading
387        log_loading.warning(
388            "IPython too old. Shell color won't be handled."
389        )
390        return
391    from scapy.config import conf
392    scapy_style = {}
393    # Overwrite colors
394    if isinstance(conf.color_theme, NoTheme):
395        shell.colors = 'nocolor'
396    elif isinstance(conf.color_theme, BrightTheme):
397        # lightbg is optimized for light backgrounds
398        shell.colors = 'lightbg'
399    elif isinstance(conf.color_theme, ColorOnBlackTheme):
400        # linux is optimised for dark backgrounds
401        shell.colors = 'linux'
402    else:
403        # default
404        shell.colors = 'neutral'
405    try:
406        get_ipython()  # type: ignore
407        # This function actually contains tons of hacks
408        color_magic = shell.magics_manager.magics["line"]["colors"]
409        color_magic(shell.colors)
410    except NameError:
411        pass
412    # Prompt Style
413    if isinstance(conf.prompt, Prompts):
414        # Set custom prompt style
415        shell.prompts_class = conf.prompt
416    else:
417        if isinstance(conf.color_theme, (FormatTheme, NoTheme)):
418            # Formatable
419            if isinstance(conf.color_theme, HTMLTheme):
420                prompt = html.escape(conf.prompt)
421            elif isinstance(conf.color_theme, LatexTheme):
422                from scapy.utils import tex_escape
423                prompt = tex_escape(conf.prompt)
424            else:
425                prompt = conf.prompt
426            prompt = conf.color_theme.prompt(prompt)
427        else:
428            # Needs to be manually set
429            prompt = str(conf.prompt)
430            scapy_style[Token.Prompt] = Color.ansi_to_pygments(
431                conf.color_theme.style_prompt
432            )
433
434        class ClassicPrompt(Prompts):
435            def in_prompt_tokens(self, cli=None):
436                # type: (Any) -> List[Tuple[Any, str]]
437                return [(Token.Prompt, prompt), ]
438
439            def out_prompt_tokens(self):
440                # type: () -> List[Tuple[Any, str]]
441                return [(Token.OutPrompt, ''), ]
442        # Apply classic prompt style
443        shell.prompts_class = ClassicPrompt
444        sys.ps1 = prompt
445    # Register scapy color style
446    shell.highlighting_style_overrides = scapy_style
447    # Apply if Live
448    try:
449        get_ipython().refresh_style()  # type: ignore
450    except NameError:
451        pass
452