audit-labs/audit-tools

A collection of scripts, queries, and other goodies you can use in an audit.

clone: git clone https://gitbay.org/audit-labs/audit-tools.git

v1.0.0: tui/app.py · raw

  1"""
  2Audit Tools — interactive terminal UI.
  3
  4Presents a platform menu, walks the user through credentials and check
  5selection, then runs the selected platform's collectors with live progress.
  6
  7Run it with:
  8
  9    python audit_tui.py
 10"""
 11
 12from typing import ClassVar
 13
 14from rich.text import Text
 15from textual import work
 16from textual.app import App, ComposeResult
 17from textual.containers import Center, Horizontal, Vertical
 18from textual.screen import Screen
 19from textual.widgets import (
 20    Button,
 21    Footer,
 22    Header,
 23    Input,
 24    Label,
 25    ProgressBar,
 26    RichLog,
 27    SelectionList,
 28    Static,
 29)
 30from textual.widgets.selection_list import Selection
 31
 32from tui import platforms
 33from tui.common import Check, ProgressEvent
 34
 35
 36class MenuScreen(Screen):
 37    """Pick a platform to audit."""
 38
 39    BINDINGS: ClassVar[list] = [("q", "app.quit", "Quit")]
 40
 41    def compose(self) -> ComposeResult:
 42        yield Header()
 43        with Center(), Vertical(id="menu-box"):
 44            yield Static("Select a platform to audit", classes="prompt")
 45            for platform in platforms.PLATFORMS:
 46                label = platform.label
 47                if not platform.enabled:
 48                    label = f"{label}  —  coming soon"
 49                yield Button(
 50                    label,
 51                    id=platform.key,
 52                    variant="primary" if platform.enabled else "default",
 53                    disabled=not platform.enabled,
 54                )
 55        yield Footer()
 56
 57    def on_mount(self) -> None:
 58        self.sub_title = "Select a platform"
 59
 60    def on_button_pressed(self, event: Button.Pressed) -> None:
 61        for platform in platforms.PLATFORMS:
 62            if event.button.id == platform.key and platform.enabled:
 63                self.app.platform = platform
 64                self.app.push_screen(ConfigScreen())
 65                return
 66
 67
 68class ConfigScreen(Screen):
 69    """Collect the connection details for the chosen platform."""
 70
 71    BINDINGS: ClassVar[list] = [("escape", "back", "Back")]
 72
 73    def compose(self) -> ComposeResult:
 74        platform = self.app.platform
 75        yield Header()
 76        with Center(), Vertical(id="form-box"):
 77            yield Static(
 78                f"{platform.label} audit — connection details", classes="prompt"
 79            )
 80            for f in platform.fields:
 81                yield Label(f.label)
 82                yield Input(
 83                    value=platforms.prefill(f),
 84                    placeholder=f.placeholder,
 85                    password=f.password,
 86                    id=f.key,
 87                )
 88            yield Static("", id="form-error", classes="error")
 89            with Horizontal(classes="buttons"):
 90                yield Button("Back", id="back")
 91                yield Button("Continue", id="continue", variant="primary")
 92        yield Footer()
 93
 94    def on_mount(self) -> None:
 95        platform = self.app.platform
 96        self.sub_title = f"{platform.label} · connection"
 97        self.query_one(f"#{platform.fields[0].key}", Input).focus()
 98
 99    def action_back(self) -> None:
100        self.app.pop_screen()
101
102    def on_button_pressed(self, event: Button.Pressed) -> None:
103        if event.button.id == "back":
104            self.app.pop_screen()
105        elif event.button.id == "continue":
106            self._submit()
107
108    def on_input_submitted(self, event: Input.Submitted) -> None:
109        self._submit()
110
111    def _submit(self) -> None:
112        platform = self.app.platform
113        settings = {}
114        missing = []
115        for f in platform.fields:
116            value = self.query_one(f"#{f.key}", Input).value.strip()
117            if not value:
118                value = f.default
119            if f.required and not value:
120                missing.append(f.label.lower())
121            settings[f.key] = value
122
123        if missing:
124            self.query_one("#form-error", Static).update(
125                f"Please provide: {', '.join(missing)}."
126            )
127            return
128
129        self.app.settings = settings
130        self.app.push_screen(ChecksScreen())
131
132
133class ChecksScreen(Screen):
134    """Choose which checks to run."""
135
136    BINDINGS: ClassVar[list] = [("escape", "back", "Back")]
137
138    def compose(self) -> ComposeResult:
139        platform = self.app.platform
140        yield Header()
141        with Center(), Vertical(id="checks-box"):
142            yield Static("Select checks to run", classes="prompt")
143            yield SelectionList(
144                *[
145                    Selection(
146                        self._prompt(c),
147                        c.key,
148                        c.key in platform.default_selection,
149                    )
150                    for c in platform.checks
151                ],
152                id="checks",
153            )
154            yield Static("", id="checks-error", classes="error")
155            with Horizontal(classes="buttons"):
156                yield Button("Back", id="back")
157                yield Button("Run audit", id="run", variant="primary")
158        yield Footer()
159
160    def on_mount(self) -> None:
161        self.sub_title = f"{self.app.platform.label} · select checks"
162        self.query_one("#checks", SelectionList).focus()
163
164    @staticmethod
165    def _prompt(check: Check) -> Text:
166        text = Text(check.label)
167        if check.note:
168            text.append(f"  ({check.note})", style="dim italic")
169        return text
170
171    def action_back(self) -> None:
172        self.app.pop_screen()
173
174    def on_button_pressed(self, event: Button.Pressed) -> None:
175        if event.button.id == "back":
176            self.app.pop_screen()
177        elif event.button.id == "run":
178            selected = list(self.query_one("#checks", SelectionList).selected)
179            if not selected:
180                self.query_one("#checks-error", Static).update(
181                    "Select at least one check."
182                )
183                return
184            self.app.selected_keys = selected
185            self.app.push_screen(RunScreen())
186
187
188class RunScreen(Screen):
189    """Run the selected checks with live progress."""
190
191    BINDINGS: ClassVar[list] = [("escape", "home", "Menu")]
192
193    def compose(self) -> ComposeResult:
194        yield Header()
195        with Vertical(id="run-box"):
196            yield Static(id="run-target", classes="prompt")
197            yield ProgressBar(id="progress", show_eta=False)
198            yield RichLog(id="log", markup=True, highlight=False, wrap=True)
199            with Horizontal(classes="buttons"):
200                yield Button("Back to menu", id="menu", disabled=True)
201                yield Button("Quit", id="quit", disabled=True, variant="primary")
202        yield Footer()
203
204    def on_mount(self) -> None:
205        platform = self.app.platform
206        settings = self.app.settings
207        keys = self.app.selected_keys
208        self.sub_title = f"{platform.label} · running"
209        self.output_dir = platform.output_dir(settings)
210        target = platform.subject(settings)
211        self.query_one("#run-target", Static).update(
212            f"Auditing [b]{target}[/]  ·  {len(keys)} checks  ·  → {self.output_dir}"
213        )
214        self._progress.update(total=len(keys), progress=0)
215        self.run_audit()
216
217    @property
218    def _progress(self) -> ProgressBar:
219        return self.query_one("#progress", ProgressBar)
220
221    @work(thread=True)
222    def run_audit(self) -> None:
223        platform = self.app.platform
224        settings = self.app.settings
225        keys = self.app.selected_keys
226        try:
227            platform.run(
228                settings,
229                self.output_dir,
230                keys,
231                lambda ev: self.app.call_from_thread(self._handle_event, ev),
232            )
233        except Exception as e:
234            self.app.call_from_thread(self._log, f"[red]Run failed:[/] {e}")
235        finally:
236            self.app.call_from_thread(self._finish)
237
238    def _log(self, markup: str) -> None:
239        self.query_one("#log", RichLog).write(markup)
240
241    def _handle_event(self, ev: ProgressEvent) -> None:
242        if ev.kind == "fetch":
243            self._log(f"[dim]· {ev.label}…[/]")
244        elif ev.kind == "start":
245            self._log(f"[cyan]▶[/] {ev.label}")
246        elif ev.kind == "done":
247            self._log(f"[green]✓[/] {ev.label} — [b]{ev.count}[/] rows")
248            self._progress.advance(1)
249        elif ev.kind == "error":
250            self._log(f"[red]✗[/] {ev.label}{ev.message}")
251            self._progress.advance(1)
252        elif ev.kind == "summary":
253            self._log("")
254            self._log(f"[bold green]Done.[/] Package written to {ev.label}")
255
256    def _finish(self) -> None:
257        self.query_one("#menu", Button).disabled = False
258        self.query_one("#quit", Button).disabled = False
259
260    def action_home(self) -> None:
261        self.app.show_menu()
262
263    def on_button_pressed(self, event: Button.Pressed) -> None:
264        if event.button.id == "menu":
265            self.app.show_menu()
266        elif event.button.id == "quit":
267            self.app.exit()
268
269
270class AuditApp(App):
271    TITLE = "Audit Tools"
272
273    CSS = """
274    Screen {
275        align: center middle;
276    }
277    #menu-box, #form-box, #checks-box {
278        width: 64;
279        height: auto;
280        padding: 1 2;
281        border: round $primary;
282    }
283    #run-box {
284        width: 90%;
285        height: 90%;
286        padding: 1 2;
287        border: round $primary;
288    }
289    .prompt {
290        text-style: bold;
291        margin-bottom: 1;
292    }
293    .error {
294        color: $error;
295        margin-top: 1;
296    }
297    Label {
298        margin-top: 1;
299    }
300    .buttons {
301        height: auto;
302        margin-top: 1;
303        align-horizontal: right;
304    }
305    .buttons Button {
306        margin-left: 2;
307    }
308    #menu-box Button {
309        width: 100%;
310        margin-top: 1;
311    }
312    #checks {
313        height: auto;
314        max-height: 14;
315    }
316    #log {
317        height: 1fr;
318        border: round $panel;
319        padding: 0 1;
320        margin-top: 1;
321    }
322    """
323
324    def on_mount(self) -> None:
325        self.platform = None
326        self.settings: dict = {}
327        self.selected_keys: list = []
328        self.push_screen(MenuScreen())
329
330    def show_menu(self) -> None:
331        """Pop back to the platform menu."""
332        while len(self.screen_stack) > 2:
333            self.pop_screen()
334
335
336def main() -> None:
337    AuditApp().run()
338
339
340if __name__ == "__main__":
341    main()