netbox-community-netbox/netbox/extras/management/commands/runscript.py

import json
import logging
import sys
import traceback
import uuid

from django.contrib.auth.models import User
from django.core.management.base import BaseCommand, CommandError
from django.db import transaction

from core.choices import JobStatusChoices
from core.models import Job
from extras.api.serializers import ScriptOutputSerializer
from extras.context_managers import change_logging
from extras.scripts import get_module_and_script
from extras.signals import clear_webhooks
from utilities.exceptions import AbortTransaction
from utilities.utils import NetBoxFakeRequest


class Command(BaseCommand):
    help = "Run a script in NetBox"

    def add_arguments(self, parser):
        parser.add_argument(
            '--loglevel',
            help="Logging Level (default: info)",
            dest='loglevel',
            default='info',
            choices=['debug', 'info', 'warning', 'error', 'critical'])
        parser.add_argument('--commit', help="Commit this script to database", action='store_true')
        parser.add_argument('--user', help="User script is running as")
        parser.add_argument('--data', help="Data as a string encapsulated JSON blob")
        parser.add_argument('script', help="Script to run")

    def handle(self, *args, **options):
        def _run_script():
            """
            Core script execution task. We capture this within a subfunction to allow for conditionally wrapping it with
            the change_logging context manager (which is bypassed if commit == False).
            """
            try:
                try:
                    with transaction.atomic():
                        script.output = script.run(data=data, commit=commit)
                        if not commit:
                            raise AbortTransaction()
                except AbortTransaction:
                    script.log_info("Database changes have been reverted automatically.")
                    clear_webhooks.send(request)
                job.data = ScriptOutputSerializer(script).data
                job.terminate()
            except Exception as e:
                stacktrace = traceback.format_exc()
                script.log_failure(
                    f"An exception occurred: `{type(e).__name__}: {e}`\n```\n{stacktrace}\n```"
                )
                script.log_info("Database changes have been reverted due to error.")
                logger.error(f"Exception raised during script execution: {e}")
                clear_webhooks.send(request)
                job.data = ScriptOutputSerializer(script).data
                job.terminate(status=JobStatusChoices.STATUS_ERRORED)

            logger.info(f"Script completed in {job.duration}")

        # Params
        script = options['script']
        loglevel = options['loglevel']
        commit = options['commit']
        try:
            data = json.loads(options['data'])
        except TypeError:
            data = {}

        module_name, script_name = script.split('.', 1)
        module, script = get_module_and_script(module_name, script_name)

        # Take user from command line if provided and exists, other
        if options['user']:
            try:
                user = User.objects.get(username=options['user'])
            except User.DoesNotExist:
                user = User.objects.filter(is_superuser=True).order_by('pk')[0]
        else:
            user = User.objects.filter(is_superuser=True).order_by('pk')[0]

        # Setup logging to Stdout
        formatter = logging.Formatter(f'[%(asctime)s][%(levelname)s] - %(message)s')
        stdouthandler = logging.StreamHandler(sys.stdout)
        stdouthandler.setLevel(logging.DEBUG)
        stdouthandler.setFormatter(formatter)

        logger = logging.getLogger(f"netbox.scripts.{script.full_name}")
        logger.addHandler(stdouthandler)

        try:
            logger.setLevel({
                'critical': logging.CRITICAL,
                'debug': logging.DEBUG,
                'error': logging.ERROR,
                'fatal': logging.FATAL,
                'info': logging.INFO,
                'warning': logging.WARNING,
            }[loglevel])
        except KeyError:
            raise CommandError(f"Invalid log level: {loglevel}")

        # Initialize the script form
        script = script()
        form = script.as_form(data, None)

        # Create the job
        job = Job.objects.create(
            object=module,
            name=script.name,
            user=User.objects.filter(is_superuser=True).order_by('pk')[0],
            job_id=uuid.uuid4()
        )

        request = NetBoxFakeRequest({
            'META': {},
            'POST': data,
            'GET': {},
            'FILES': {},
            'user': user,
            'path': '',
            'id': job.job_id
        })

        if form.is_valid():
            job.status = JobStatusChoices.STATUS_RUNNING
            job.save()

            logger.info(f"Running script (commit={commit})")
            script.request = request

            # Execute the script. If commit is True, wrap it with the change_logging context manager to ensure we process
            # change logging, webhooks, etc.
            with change_logging(request):
                _run_script()
        else:
            logger.error('Data is not valid:')
            for field, errors in form.errors.get_json_data().items():
                for error in errors:
                    logger.error(f'\t{field}: {error.get("message")}')
            job.status = JobStatusChoices.STATUS_ERRORED
            job.save()