1## This file is part of Scapy
2## See http://www.secdev.org/projects/scapy for more informations
3## Copyright (C) Philippe Biondi <phil@secdev.org>
4## This program is published under a GPLv2 license
5
6"""
7Color themes for the interactive console.
8"""
9
10##################
11## Color themes ##
12##################
13
14class ColorTable:
15    colors = { # Format: (ansi, pygments)
16        "normal": ("\033[0m", "noinherit"),
17        "black": ("\033[30m", "#ansiblack"),
18        "red": ("\033[31m", "#ansired"),
19        "green": ("\033[32m", "#ansigreen"),
20        "yellow": ("\033[33m", "#ansiyellow"),
21        "blue": ("\033[34m", "#ansiblue"),
22        "purple": ("\033[35m", "#ansipurple"),
23        "cyan": ("\033[36m", "#ansicyan"),
24        "grey": ("\033[37m", "#ansigrey"),
25
26        "bold": ("\033[1m", "bold"),
27        "uline": ("\033[4m", "underline"),
28        "blink": ("\033[5m", ""),
29        "invert": ("\033[7m", ""),
30        }
31
32    def __repr__(self):
33        return "<ColorTable>"
34
35    def __getattr__(self, attr):
36        return self.colors.get(attr, [""])[0]
37
38    def ansi_to_pygments(self, x): # Transform ansi encoded text to Pygments text
39        inv_map = {v[0]: v[1] for k, v in self.colors.items()}
40        for k, v in inv_map.items():
41            x = x.replace(k, " "+v)
42        return x.strip()
43
44Color = ColorTable()
45
46def create_styler(fmt=None, before="", after="", fmt2="%s"):
47    def do_style(val, fmt=fmt, before=before, after=after, fmt2=fmt2):
48        if fmt is None:
49            if not isinstance(val, str):
50                val = str(val)
51        else:
52            val = fmt % val
53        return fmt2 % (before+val+after)
54    return do_style
55
56class ColorTheme:
57    def __repr__(self):
58        return "<%s>" % self.__class__.__name__
59    def __reduce__(self):
60        return (self.__class__, (), ())
61    def __getattr__(self, attr):
62        if attr in ["__getstate__", "__setstate__", "__getinitargs__",
63                    "__reduce_ex__"]:
64            raise AttributeError()
65        return create_styler()
66
67
68class NoTheme(ColorTheme):
69    pass
70
71
72class AnsiColorTheme(ColorTheme):
73    def __getattr__(self, attr):
74        if attr.startswith("__"):
75            raise AttributeError(attr)
76        s = "style_%s" % attr
77        if s in self.__class__.__dict__:
78            before = getattr(self, s)
79            after = self.style_normal
80        else:
81            before = after = ""
82
83        return create_styler(before=before, after=after)
84
85
86    style_normal = ""
87    style_prompt = ""
88    style_punct = ""
89    style_id = ""
90    style_not_printable = ""
91    style_layer_name = ""
92    style_field_name = ""
93    style_field_value = ""
94    style_emph_field_name = ""
95    style_emph_field_value = ""
96    style_packetlist_name = ""
97    style_packetlist_proto = ""
98    style_packetlist_value = ""
99    style_fail = ""
100    style_success = ""
101    style_odd = ""
102    style_even = ""
103    style_opening = ""
104    style_active = ""
105    style_closed = ""
106    style_left = ""
107    style_right = ""
108    style_logo = ""
109
110class BlackAndWhite(AnsiColorTheme):
111    pass
112
113class DefaultTheme(AnsiColorTheme):
114    style_normal = Color.normal
115    style_prompt = Color.blue+Color.bold
116    style_punct = Color.normal
117    style_id = Color.blue+Color.bold
118    style_not_printable = Color.grey
119    style_layer_name = Color.red+Color.bold
120    style_field_name = Color.blue
121    style_field_value = Color.purple
122    style_emph_field_name = Color.blue+Color.uline+Color.bold
123    style_emph_field_value = Color.purple+Color.uline+Color.bold
124    style_packetlist_name = Color.red+Color.bold
125    style_packetlist_proto = Color.blue
126    style_packetlist_value = Color.purple
127    style_fail = Color.red+Color.bold
128    style_success = Color.blue+Color.bold
129    style_even = Color.black+Color.bold
130    style_odd = Color.black
131    style_opening = Color.yellow
132    style_active = Color.black
133    style_closed = Color.grey
134    style_left = Color.blue+Color.invert
135    style_right = Color.red+Color.invert
136    style_logo = Color.green+Color.bold
137
138class BrightTheme(AnsiColorTheme):
139    style_normal = Color.normal
140    style_punct = Color.normal
141    style_id = Color.yellow+Color.bold
142    style_layer_name = Color.red+Color.bold
143    style_field_name = Color.yellow+Color.bold
144    style_field_value = Color.purple+Color.bold
145    style_emph_field_name = Color.yellow+Color.bold
146    style_emph_field_value = Color.green+Color.bold
147    style_packetlist_name = Color.red+Color.bold
148    style_packetlist_proto = Color.yellow+Color.bold
149    style_packetlist_value = Color.purple+Color.bold
150    style_fail = Color.red+Color.bold
151    style_success = Color.blue+Color.bold
152    style_even = Color.black+Color.bold
153    style_odd = Color.black
154    style_left = Color.cyan+Color.invert
155    style_right = Color.purple+Color.invert
156    style_logo = Color.green+Color.bold
157
158
159class RastaTheme(AnsiColorTheme):
160    style_normal = Color.normal+Color.green+Color.bold
161    style_prompt = Color.yellow+Color.bold
162    style_punct = Color.red
163    style_id = Color.green+Color.bold
164    style_not_printable = Color.green
165    style_layer_name = Color.red+Color.bold
166    style_field_name = Color.yellow+Color.bold
167    style_field_value = Color.green+Color.bold
168    style_emph_field_name = Color.green
169    style_emph_field_value = Color.green
170    style_packetlist_name = Color.red+Color.bold
171    style_packetlist_proto = Color.yellow+Color.bold
172    style_packetlist_value = Color.green+Color.bold
173    style_fail = Color.red
174    style_success = Color.red+Color.bold
175    style_even = Color.yellow
176    style_odd = Color.green
177    style_left = Color.yellow+Color.invert
178    style_right = Color.red+Color.invert
179    style_logo = Color.green+Color.bold
180
181
182class ColorOnBlackTheme(AnsiColorTheme):
183    """Color theme for black backgrounds"""
184    style_normal = Color.normal
185    style_prompt = Color.green+Color.bold
186    style_punct = Color.normal
187    style_id = Color.green
188    style_not_printable = Color.black+Color.bold
189    style_layer_name = Color.yellow+Color.bold
190    style_field_name = Color.cyan
191    style_field_value = Color.purple+Color.bold
192    style_emph_field_name = Color.cyan+Color.bold
193    style_emph_field_value = Color.red+Color.bold
194    style_packetlist_name = Color.black+Color.bold
195    style_packetlist_proto = Color.yellow+Color.bold
196    style_packetlist_value = Color.purple+Color.bold
197    style_fail = Color.red+Color.bold
198    style_success = Color.green
199    style_even = Color.black+Color.bold
200    style_odd = Color.grey
201    style_opening = Color.yellow
202    style_active = Color.grey+Color.bold
203    style_closed = Color.black+Color.bold
204    style_left = Color.cyan+Color.bold
205    style_right = Color.red+Color.bold
206    style_logo = Color.green+Color.bold
207
208
209class FormatTheme(ColorTheme):
210    def __getattr__(self, attr):
211        if attr.startswith("__"):
212            raise AttributeError(attr)
213        colfmt = self.__class__.__dict__.get("style_%s" % attr, "%s")
214        return create_styler(fmt2 = colfmt)
215
216class LatexTheme(FormatTheme):
217    style_prompt = r"\textcolor{blue}{%s}"
218    style_not_printable = r"\textcolor{gray}{%s}"
219    style_layer_name = r"\textcolor{red}{\bf %s}"
220    style_field_name = r"\textcolor{blue}{%s}"
221    style_field_value = r"\textcolor{purple}{%s}"
222    style_emph_field_name = r"\textcolor{blue}{\underline{%s}}" #ul
223    style_emph_field_value = r"\textcolor{purple}{\underline{%s}}" #ul
224    style_packetlist_name = r"\textcolor{red}{\bf %s}"
225    style_packetlist_proto = r"\textcolor{blue}{%s}"
226    style_packetlist_value = r"\textcolor{purple}{%s}"
227    style_fail = r"\textcolor{red}{\bf %s}"
228    style_success = r"\textcolor{blue}{\bf %s}"
229    style_left = r"\textcolor{blue}{%s}"
230    style_right = r"\textcolor{red}{%s}"
231#    style_even = r"}{\bf "
232#    style_odd = ""
233    style_logo = r"\textcolor{green}{\bf %s}"
234
235class LatexTheme2(FormatTheme):
236    style_prompt = r"@`@textcolor@[@blue@]@@[@%s@]@"
237    style_not_printable = r"@`@textcolor@[@gray@]@@[@%s@]@"
238    style_layer_name = r"@`@textcolor@[@red@]@@[@@`@bfseries@[@@]@%s@]@"
239    style_field_name = r"@`@textcolor@[@blue@]@@[@%s@]@"
240    style_field_value = r"@`@textcolor@[@purple@]@@[@%s@]@"
241    style_emph_field_name = r"@`@textcolor@[@blue@]@@[@@`@underline@[@%s@]@@]@"
242    style_emph_field_value = r"@`@textcolor@[@purple@]@@[@@`@underline@[@%s@]@@]@"
243    style_packetlist_name = r"@`@textcolor@[@red@]@@[@@`@bfseries@[@@]@%s@]@"
244    style_packetlist_proto = r"@`@textcolor@[@blue@]@@[@%s@]@"
245    style_packetlist_value = r"@`@textcolor@[@purple@]@@[@%s@]@"
246    style_fail = r"@`@textcolor@[@red@]@@[@@`@bfseries@[@@]@%s@]@"
247    style_success = r"@`@textcolor@[@blue@]@@[@@`@bfserices@[@@]@%s@]@"
248    style_even = r"@`@textcolor@[@gray@]@@[@@`@bfseries@[@@]@%s@]@"
249#    style_odd = r"@`@textcolor@[@black@]@@[@@`@bfseries@[@@]@%s@]@"
250    style_left = r"@`@textcolor@[@blue@]@@[@%s@]@"
251    style_right = r"@`@textcolor@[@red@]@@[@%s@]@"
252    style_logo = r"@`@textcolor@[@green@]@@[@@`@bfseries@[@@]@%s@]@"
253
254class HTMLTheme(FormatTheme):
255    style_prompt = "<span class=prompt>%s</span>"
256    style_not_printable = "<span class=not_printable>%s</span>"
257    style_layer_name = "<span class=layer_name>%s</span>"
258    style_field_name = "<span class=field_name>%s</span>"
259    style_field_value = "<span class=field_value>%s</span>"
260    style_emph_field_name = "<span class=emph_field_name>%s</span>"
261    style_emph_field_value = "<span class=emph_field_value>%s</span>"
262    style_packetlist_name = "<span class=packetlist_name>%s</span>"
263    style_packetlist_proto = "<span class=packetlist_proto>%s</span>"
264    style_packetlist_value = "<span class=packetlist_value>%s</span>"
265    style_fail = "<span class=fail>%s</span>"
266    style_success = "<span class=success>%s</span>"
267    style_even = "<span class=even>%s</span>"
268    style_odd = "<span class=odd>%s</span>"
269    style_left = "<span class=left>%s</span>"
270    style_right = "<span class=right>%s</span>"
271
272class HTMLTheme2(HTMLTheme):
273    style_prompt = "#[#span class=prompt#]#%s#[#/span#]#"
274    style_not_printable = "#[#span class=not_printable#]#%s#[#/span#]#"
275    style_layer_name = "#[#span class=layer_name#]#%s#[#/span#]#"
276    style_field_name = "#[#span class=field_name#]#%s#[#/span#]#"
277    style_field_value = "#[#span class=field_value#]#%s#[#/span#]#"
278    style_emph_field_name = "#[#span class=emph_field_name#]#%s#[#/span#]#"
279    style_emph_field_value = "#[#span class=emph_field_value#]#%s#[#/span#]#"
280    style_packetlist_name = "#[#span class=packetlist_name#]#%s#[#/span#]#"
281    style_packetlist_proto = "#[#span class=packetlist_proto#]#%s#[#/span#]#"
282    style_packetlist_value = "#[#span class=packetlist_value#]#%s#[#/span#]#"
283    style_fail = "#[#span class=fail#]#%s#[#/span#]#"
284    style_success = "#[#span class=success#]#%s#[#/span#]#"
285    style_even = "#[#span class=even#]#%s#[#/span#]#"
286    style_odd = "#[#span class=odd#]#%s#[#/span#]#"
287    style_left = "#[#span class=left#]#%s#[#/span#]#"
288    style_right = "#[#span class=right#]#%s#[#/span#]#"
289
290
291def apply_ipython_style(shell):
292    """Updates the specified IPython console shell with
293    the conf.color_theme scapy theme."""
294    try:
295        from IPython.terminal.prompts import Prompts, Token
296    except:
297        from scapy.error import log_loading
298        log_loading.warning(
299            "IPython too old. Shell color won't be handled."
300            )
301        return
302    from scapy.config import conf
303    if isinstance(conf.prompt, Prompts):
304        shell.prompts_class = conf.prompt # Set custom prompt style
305    else:
306        class ClassicPrompt(Prompts):
307            def in_prompt_tokens(self, cli=None):
308               return [(Token.Prompt, str(conf.prompt)),]
309            def out_prompt_tokens(self):
310               return [(Token.OutPrompt, ''),]
311        shell.prompts_class=ClassicPrompt # Apply classic prompt style
312    shell.highlighting_style_overrides = { # Register and apply scapy color style
313        Token.Prompt: Color.ansi_to_pygments(conf.color_theme.style_prompt),
314    }
315    try:
316        get_ipython().refresh_style()
317    except NameError:
318        pass
319