diff --git a/base_requirements.txt b/base_requirements.txt index 7ba631bb9..8d74ede6c 100644 --- a/base_requirements.txt +++ b/base_requirements.txt @@ -1,3 +1,7 @@ +# Shell text coloring +# https://github.com/tartley/colorama/blob/master/CHANGELOG.rst +colorama + # The Python web framework on which NetBox is built # https://docs.djangoproject.com/en/stable/releases/ Django==5.2.* diff --git a/netbox/core/management/commands/nbshell.py b/netbox/core/management/commands/nbshell.py index 8f729d10a..ef7ae191e 100644 --- a/netbox/core/management/commands/nbshell.py +++ b/netbox/core/management/commands/nbshell.py @@ -1,29 +1,48 @@ import code import platform -import sys +from collections import defaultdict +from types import SimpleNamespace +from colorama import Fore, Style from django import get_version from django.apps import apps from django.conf import settings from django.core.management.base import BaseCommand +from django.utils.module_loading import import_string -from core.models import ObjectType -from users.models import User +from netbox.constants import CORE_APPS +from netbox.plugins.utils import get_installed_plugins -APPS = ('circuits', 'core', 'dcim', 'extras', 'ipam', 'tenancy', 'users', 'virtualization', 'vpn', 'wireless') -EXCLUDE_MODELS = ( - 'extras.branch', - 'extras.stagedchange', -) -BANNER_TEXT = """### NetBox interactive shell ({node}) -### Python {python} | Django {django} | NetBox {netbox} -### lsmodels() will show available models. Use help() for more info.""".format( - node=platform.node(), - python=platform.python_version(), - django=get_version(), - netbox=settings.RELEASE.name -) +def color(color: str, text: str): + return getattr(Fore, color.upper()) + text + Style.RESET_ALL + + +def bright(text: str): + return Style.BRIGHT + text + Style.RESET_ALL + + +def get_models(app_config): + """ + Return a list of all non-private models within an app. + """ + return [ + model for model in app_config.get_models() + if not getattr(model, '_netbox_private', False) + ] + + +def get_constants(app_config): + """ + Return a dictionary mapping of all constants defined within an app. + """ + try: + constants = import_string(f'{app_config.name}.constants') + except ImportError: + return {} + return { + name: value for name, value in vars(constants).items() + } class Command(BaseCommand): @@ -36,47 +55,88 @@ class Command(BaseCommand): help='Python code to execute (instead of starting an interactive shell)', ) - def _lsmodels(self): - for app, models in self.django_models.items(): - app_name = apps.get_app_config(app).verbose_name + def _lsapps(self): + for app_label in self.django_models.keys(): + app_name = apps.get_app_config(app_label).verbose_name + print(f'{app_label} - {app_name}') + + def _lsmodels(self, app_label=None): + """ + Return a list of all models within each app. + + Args: + app_label: The name of a specific app + """ + if app_label: + if app_label not in self.django_models: + print(f"No models listed for {app_label}") + return + app_labels = [app_label] + else: + app_labels = self.django_models.keys() # All apps + + for app_label in app_labels: + app_name = apps.get_app_config(app_label).verbose_name print(f'{app_name}:') - for m in models: + for m in self.django_models[app_label]: print(f' {m}') def get_namespace(self): - namespace = {} + namespace = defaultdict(SimpleNamespace) - # Gather Django models and constants from each app - for app in APPS: - models = [] + # Iterate through all core apps & plugins to compile namespace of models and constants + for app_name in [*CORE_APPS, *get_installed_plugins().keys()]: + app_config = apps.get_app_config(app_name) - # Load models from each app - for model in apps.get_app_config(app).get_models(): - app_label = model._meta.app_label - model_name = model._meta.model_name - if f'{app_label}.{model_name}' not in EXCLUDE_MODELS: - namespace[model.__name__] = model - models.append(model.__name__) - self.django_models[app] = sorted(models) + # Populate models + if models := get_models(app_config): + for model in models: + setattr(namespace[app_name], model.__name__, model) + self.django_models[app_name] = sorted([ + model.__name__ for model in models + ]) - # Constants - try: - app_constants = sys.modules[f'{app}.constants'] - for name in dir(app_constants): - namespace[name] = getattr(app_constants, name) - except KeyError: - pass + # Populate constants + for const_name, const_value in get_constants(app_config).items(): + setattr(namespace[app_name], const_name, const_value) - # Additional objects to include - namespace['ObjectType'] = ObjectType - namespace['User'] = User - - # Load convenience commands - namespace.update({ + return { + **namespace, + 'lsapps': self._lsapps, 'lsmodels': self._lsmodels, - }) + } - return namespace + @staticmethod + def get_banner_text(): + lines = [ + '{title} ({hostname})'.format( + title=bright('NetBox interactive shell'), + hostname=platform.node(), + ), + '{python} | {django} | {netbox}'.format( + python=color('green', f'Python v{platform.python_version()}'), + django=color('green', f'Django v{get_version()}'), + netbox=color('green', settings.RELEASE.name), + ), + ] + + if installed_plugins := get_installed_plugins(): + plugin_list = ', '.join([ + color('cyan', f'{name} v{version}') for name, version in installed_plugins.items() + ]) + lines.append( + 'Plugins: {plugin_list}'.format( + plugin_list=plugin_list + ) + ) + + lines.append( + 'lsapps() & lsmodels() will show available models. Use help() for more info.' + ) + + return '\n'.join([ + f'### {line}' for line in lines + ]) def handle(self, **options): namespace = self.get_namespace() @@ -97,5 +157,4 @@ class Command(BaseCommand): readline.parse_and_bind('tab: complete') # Run interactive shell - shell = code.interact(banner=BANNER_TEXT, local=namespace) - return shell + return code.interact(banner=self.get_banner_text(), local=namespace) diff --git a/requirements.txt b/requirements.txt index 28481d445..80decd027 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +colorama==0.4.6 Django==5.2.4 django-cors-headers==4.7.0 django-debug-toolbar==5.2.0