Skip to content

Commit e5a0303

Browse files
committed
Set screen formatter colours based on terminal BG
Two-stage background detection: $COLORFGBG followed by OSC11. Fallback to dark if colour can't be detected using the methods above. Can be overriden by $BANDIT_LIGHT_BG environment variable.
1 parent 61d1667 commit e5a0303

File tree

1 file changed

+157
-7
lines changed

1 file changed

+157
-7
lines changed

bandit/formatters/screen.py

Lines changed: 157 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,20 @@
3131
.. versionchanged:: 1.7.3
3232
New field `CWE` added to output
3333
34+
.. versionchanged:: 1.8.5
35+
Automatic colours configuration with optional override via
36+
"BANDIT_LIGHT_BG" environment variable.
37+
3438
"""
39+
3540
import datetime
3641
import logging
42+
import os
43+
import select
3744
import sys
45+
import termios
46+
import time
47+
import tty
3848

3949
from bandit.core import constants
4050
from bandit.core import docs_utils
@@ -57,13 +67,153 @@
5767

5868
LOG = logging.getLogger(__name__)
5969

60-
COLOR = {
61-
"DEFAULT": "\033[0m",
62-
"HEADER": "\033[95m",
63-
"LOW": "\033[94m",
64-
"MEDIUM": "\033[93m",
65-
"HIGH": "\033[91m",
66-
}
70+
71+
def term_detect_bg() -> bool | None:
72+
"""Detects if terminal is using dark BG.
73+
74+
Returns:
75+
True - Light
76+
False - Dark
77+
None - Undetermined
78+
"""
79+
colorfgbg = os.environ.get("COLORFGBG")
80+
if colorfgbg and ";" in colorfgbg:
81+
try:
82+
parts = colorfgbg.split(";")
83+
bg_color = int(parts[-1])
84+
# Ref. https://github.com/rocky/shell-term-background
85+
if bg_color in {0, 1, 2, 3, 4, 5, 6, 8}:
86+
return False
87+
elif bg_color in {7, 9, 10, 11, 12, 13, 14, 15}:
88+
return True
89+
except (ValueError, IndexError):
90+
pass
91+
if sys.stdin.isatty():
92+
try:
93+
result = term_get_osc()
94+
if result is not None:
95+
return result
96+
except Exception:
97+
pass
98+
if os.environ.get("BANDIT_LIGHT_BG", "").lower() in (
99+
"light", "bright", "white", "1", "true", "yes"
100+
):
101+
return True
102+
103+
return None
104+
105+
106+
def term_get_osc() -> bool | None:
107+
"""Query terminal BG colour using OSC11.
108+
109+
Returns:
110+
True - Light
111+
False - Dark
112+
None - Undetermined
113+
"""
114+
if not sys.stdin.isatty():
115+
return None
116+
117+
old_settings = None
118+
119+
try:
120+
old_settings = termios.tcgetattr(sys.stdin)
121+
122+
_ = tty.setraw(sys.stdin.fileno())
123+
_ = sys.stdout.write("\x1b]11;?\x1b\\") # ESC\
124+
_ = sys.stdout.flush()
125+
126+
ready, _, _ = select.select([sys.stdin], [], [], 0.2)
127+
if not ready:
128+
return None # Bail out, this term is cursed
129+
130+
response = ""
131+
start_time = time.time()
132+
while time.time() - start_time < 0.5:
133+
ready, _, _ = select.select([sys.stdin], [], [], 0.01)
134+
if not ready:
135+
break
136+
char = sys.stdin.read(1)
137+
response += char
138+
# Break on ESC\, BEL or sufficient data
139+
if response.endswith('\x1b\\') or response.endswith('\x07') or len(response) > 50:
140+
break
141+
# Bail out if ESC isn't followed by ]
142+
if len(response) >= 2 and response.startswith('\x1b') and not response.startswith('\x1b]'):
143+
return None
144+
145+
if response.startswith('\x1b]11;rgb:'):
146+
try:
147+
rgb_start = response.find("rgb:")
148+
rgb_part = response[rgb_start + 4:]
149+
# Find terminator
150+
for term in ['\x1b\\', '\x07']:
151+
if term in rgb_part:
152+
rgb_part = rgb_part[:rgb_part.find(term)]
153+
break
154+
155+
r, g, b = rgb_part.split("/")[:3]
156+
157+
# HEX -> DEC
158+
r_val = int(r[:4], 16) if len(r) >= 4 else int(r, 16)
159+
g_val = int(g[:4], 16) if len(g) >= 4 else int(g, 16)
160+
b_val = int(b[:4], 16) if len(b) >= 4 else int(b, 16)
161+
162+
# 16b -> 8b
163+
if r_val > 255:
164+
r_val = r_val >> 8
165+
if g_val > 255:
166+
g_val = g_val >> 8
167+
if b_val > 255:
168+
b_val = b_val >> 8
169+
170+
# BT601
171+
lum = 0.299 * r_val + 0.587 * g_val + 0.114 * b_val
172+
return lum > 128 # Light if luma > 50% grey
173+
except (ValueError, IndexError):
174+
pass
175+
else:
176+
return None
177+
178+
except Exception:
179+
pass
180+
finally:
181+
try:
182+
if old_settings:
183+
termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)
184+
except Exception:
185+
pass
186+
187+
return None
188+
189+
190+
def term_serve_colourscheme() -> dict[str, str]:
191+
"""Appropriate colour scheme based on the detected background.
192+
193+
Returns:
194+
Dictionary with colour codes
195+
"""
196+
light = term_detect_bg()
197+
198+
if light:
199+
return {
200+
"DEFAULT": "\033[0m",
201+
"HEADER": "\033[1;34m", # Dark blue
202+
"LOW": "\033[1;32m", # Dark green
203+
"MEDIUM": "\033[1;35m", # Dark magenta
204+
"HIGH": "\033[1;31m", # Dark red
205+
}
206+
else:
207+
return {
208+
"DEFAULT": "\033[0m",
209+
"HEADER": "\033[1;96m", # Bright cyan
210+
"LOW": "\033[1;92m", # Bright green
211+
"MEDIUM": "\033[1;93m", # Bright yellow
212+
"HIGH": "\033[1;91m", # Bright red
213+
}
214+
215+
216+
COLOR = term_serve_colourscheme()
67217

68218

69219
def header(text, *args):

0 commit comments

Comments
 (0)