1import code
2import platform
3import sys
4
5from django import get_version
6from django.apps import apps
7from django.conf import settings
8from django.contrib.auth.models import User
9from django.contrib.contenttypes.models import ContentType
10from django.core.management.base import BaseCommand
11
12APPS = ['circuits', 'dcim', 'extras', 'ipam', 'tenancy', 'users', 'virtualization']
13
14BANNER_TEXT = """### NetBox interactive shell ({node})
15### Python {python} | Django {django} | NetBox {netbox}
16### lsmodels() will show available models. Use help(<model>) for more info.""".format(
17    node=platform.node(),
18    python=platform.python_version(),
19    django=get_version(),
20    netbox=settings.VERSION
21)
22
23
24class Command(BaseCommand):
25    help = "Start the Django shell with all NetBox models already imported"
26    django_models = {}
27
28    def add_arguments(self, parser):
29        parser.add_argument(
30            '-c', '--command',
31            help='Python code to execute (instead of starting an interactive shell)',
32        )
33
34    def _lsmodels(self):
35        for app, models in self.django_models.items():
36            app_name = apps.get_app_config(app).verbose_name
37            print(f'{app_name}:')
38            for m in models:
39                print(f'  {m}')
40
41    def get_namespace(self):
42        namespace = {}
43
44        # Gather Django models and constants from each app
45        for app in APPS:
46            self.django_models[app] = []
47
48            # Load models from each app
49            for model in apps.get_app_config(app).get_models():
50                namespace[model.__name__] = model
51                self.django_models[app].append(model.__name__)
52
53            # Constants
54            try:
55                app_constants = sys.modules[f'{app}.constants']
56                for name in dir(app_constants):
57                    namespace[name] = getattr(app_constants, name)
58            except KeyError:
59                pass
60
61        # Additional objects to include
62        namespace['ContentType'] = ContentType
63        namespace['User'] = User
64
65        # Load convenience commands
66        namespace.update({
67            'lsmodels': self._lsmodels,
68        })
69
70        return namespace
71
72    def handle(self, **options):
73        # If Python code has been passed, execute it and exit.
74        if options['command']:
75            exec(options['command'], self.get_namespace())
76            return
77
78        shell = code.interact(banner=BANNER_TEXT, local=self.get_namespace())
79        return shell
80