Coverage for extras/management/commands/runscript.py: 0%

1import json 

2import logging 

3import sys 

4import uuid 

5 

6from django.core.exceptions import ValidationError 

7from django.core.management.base import BaseCommand, CommandError 

8 

9from extras.jobs import ScriptJob 

10from extras.scripts import EXEC_PARAM_FIELDS, get_module_and_script 

11from users.models import User 

12from utilities.request import NetBoxFakeRequest 

13 

14 

15class Command(BaseCommand): 

16 help = "Run a script in NetBox" 

17 

18 def add_arguments(self, parser): 

19 parser.add_argument( 

20 '--loglevel', 

21 help="Logging Level (default: info)", 

22 dest='loglevel', 

23 default='info', 

24 choices=['debug', 'info', 'warning', 'error', 'critical']) 

25 parser.add_argument('--commit', help="Commit this script to database", action='store_true') 

26 parser.add_argument('--user', help="User script is running as") 

27 parser.add_argument('--data', help="Data as a string encapsulated JSON blob") 

28 parser.add_argument('script', help="Script to run") 

29 

30 def handle(self, *args, **options): 

31 # Params 

32 script = options['script'] 

33 loglevel = options['loglevel'] 

34 commit = options['commit'] 

35 

36 try: 

37 data = json.loads(options['data']) 

38 except TypeError: 

39 data = {} 

40 

41 module_name, script_name = script.split('.', 1) 

42 script_obj = get_module_and_script(module_name, script_name)[1] 

43 script = script_obj.python_class 

44 

45 # Take user from command line if provided and exists, other 

46 if options['user']: 

47 try: 

48 user = User.objects.get(username=options['user']) 

49 except User.DoesNotExist: 

50 user = User.objects.filter(is_superuser=True).order_by('pk')[0] 

51 else: 

52 user = User.objects.filter(is_superuser=True).order_by('pk')[0] 

53 

54 # Setup logging to Stdout 

55 formatter = logging.Formatter('[%(asctime)s][%(levelname)s] - %(message)s') 

56 stdouthandler = logging.StreamHandler(sys.stdout) 

57 stdouthandler.setLevel(logging.DEBUG) 

58 stdouthandler.setFormatter(formatter) 

59 

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

61 logger.addHandler(stdouthandler) 

62 

63 try: 

64 logger.setLevel({ 

65 'critical': logging.CRITICAL, 

66 'debug': logging.DEBUG, 

67 'error': logging.ERROR, 

68 'fatal': logging.FATAL, 

69 'info': logging.INFO, 

70 'warning': logging.WARNING, 

71 }[loglevel]) 

72 except KeyError: # pragma: no cover 

73 raise CommandError(f"Invalid log level: {loglevel}") 

74 

75 # Initialize the script form 

76 script = script() 

77 form = script.as_form(data, None) 

78 if not form.is_valid(): 

79 logger.error('Data is not valid:') 

80 for field, errors in form.errors.get_json_data().items(): 

81 for error in errors: 

82 logger.error(f'\t{field}: {error.get("message")}') 

83 raise CommandError() 

84 

85 # Remove exec-parameter fields from ScriptForm before passing data to the script 

86 cleaned_data = form.cleaned_data.copy() 

87 notifications = cleaned_data.pop('_notifications') 

88 for key in EXEC_PARAM_FIELDS: 

89 cleaned_data.pop(key, None) 

90 

91 # Execute the script. 

92 try: 

93 job = ScriptJob.enqueue( 

94 instance=script_obj, 

95 user=user, 

96 immediate=True, 

97 data=cleaned_data, 

98 notifications=notifications, 

99 request=NetBoxFakeRequest({ 

100 'META': {}, 

101 'COOKIES': {}, 

102 'POST': data, 

103 'GET': {}, 

104 'FILES': {}, 

105 'user': user, 

106 'method': 'POST', 

107 'path': '', 

108 'id': uuid.uuid4() 

109 }), 

110 commit=commit, 

111 ) 

112 except ValidationError as e: 

113 # The script's Meta configuration is invalid (see #22872). Report it as a clean command error rather than 

114 # an unhandled traceback. 

115 raise CommandError('; '.join(e.messages)) 

116 

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