Coverage for extras/models/mixins.py: 29%

1import importlib.abc 

2import importlib.util 

3import logging 

4import os 

5import sys 

6from collections import defaultdict 

7 

8from django.core.exceptions import ValidationError 

9from django.core.files.storage import storages 

10from django.db import models 

11from django.http import HttpResponse 

12from django.utils.http import content_disposition_header 

13from django.utils.module_loading import import_string 

14from django.utils.translation import gettext_lazy as _ 

15 

16from core.models import ObjectType 

17from extras.constants import DEFAULT_MIME_TYPE, JINJA_ENV_PARAMS_ALLOWED, SCRIPT_MODULE_NAME_PREFIX 

18from extras.utils import filename_from_model, filename_from_object 

19from utilities.jinja2 import render_jinja2 

20 

21__all__ = ( 

22 'PythonModuleMixin', 

23 'RenderTemplateMixin', 

24) 

25 

26logger = logging.getLogger(__name__) 

27 

28 

29class CustomStoragesLoader(importlib.abc.Loader): 

30 """ 

31 Custom loader for exec_module to use django-storages instead of the file system. 

32 """ 

33 def __init__(self, filename): 

34 self.filename = filename 

35 

36 def create_module(self, spec): 

37 return None # Use default module creation 

38 

39 def exec_module(self, module): 

40 with storages["scripts"].open(self.filename, 'rb') as f: 

41 code = f.read() 

42 exec(code, module.__dict__) 

43 

44 

45class PythonModuleMixin: 

46 

47 def get_jobs(self, name): 

48 """ 

49 Returns a list of Jobs associated with this specific script or report module 

50 :param name: The class name of the script or report 

51 :return: List of Jobs associated with this 

52 """ 

53 return self.jobs.filter( 

54 name=name 

55 ) 

56 

57 @property 

58 def path(self): 

59 return os.path.splitext(self.file_path)[0] 

60 

61 @property 

62 def python_name(self): 

63 path, filename = os.path.split(self.full_path) 

64 name = os.path.splitext(filename)[0] 

65 if name == '__init__': 

66 # File is a package 

67 return os.path.basename(path) 

68 return name 

69 

70 def get_module(self): 

71 """ 

72 Load the module using importlib, but use a custom loader to use django-storages 

73 instead of the file system. 

74 """ 

75 # Load the module under a namespaced name rather than its bare name. Using the bare name 

76 # (e.g. "circuits") would replace the like-named core app package in sys.modules, breaking 

77 # app and migration graph resolution. 

78 module_name = f'{SCRIPT_MODULE_NAME_PREFIX}{self.python_name}' 

79 spec = importlib.util.spec_from_file_location(module_name, self.name) 

80 if spec is None: 

81 raise ModuleNotFoundError(f"Could not find module: {self.python_name}") 

82 loader = CustomStoragesLoader(self.name) 

83 module = importlib.util.module_from_spec(spec) 

84 sys.modules[module_name] = module 

85 loader.exec_module(module) 

86 

87 return module 

88 

89 

90class RenderTemplateMixin(models.Model): 

91 """ 

92 Enables support for rendering templates. 

93 """ 

94 template_code = models.TextField( 

95 verbose_name=_('template code'), 

96 help_text=_('Jinja template code.') 

97 ) 

98 environment_params = models.JSONField( 

99 verbose_name=_('environment parameters'), 

100 blank=True, 

101 null=True, 

102 default=dict, 

103 help_text=_( 

104 'Any <a href="{url}">additional parameters</a> to pass when constructing the Jinja environment' 

105 ).format(url='https://jinja.palletsprojects.com/en/stable/api/#jinja2.Environment') 

106 ) 

107 mime_type = models.CharField( 

108 max_length=50, 

109 blank=True, 

110 verbose_name=_('MIME type'), 

111 help_text=_('Defaults to <code>{default}</code>').format(default=DEFAULT_MIME_TYPE), 

112 ) 

113 file_name = models.CharField( 

114 max_length=200, 

115 blank=True, 

116 help_text=_('Filename to give to the rendered export file') 

117 ) 

118 file_extension = models.CharField( 

119 verbose_name=_('file extension'), 

120 max_length=15, 

121 blank=True, 

122 help_text=_('Extension to append to the rendered filename') 

123 ) 

124 as_attachment = models.BooleanField( 

125 verbose_name=_('as attachment'), 

126 default=True, 

127 help_text=_("Download file as attachment") 

128 ) 

129 

130 class Meta: 

131 abstract = True 

132 

133 def get_context(self, context=None, queryset=None): 

134 from django.apps import apps as django_apps 

135 

136 from netbox.plugins import PluginConfig 

137 

138 _context = defaultdict(dict) 

139 

140 # Populate all public models for reference within the template 

141 for object_type in ObjectType.objects.public(): 

142 if model := object_type.model_class(): 

143 _context[object_type.app_label][model.__name__] = model 

144 

145 # Allow plugins to inject additional context (e.g. friendly-named namespaces) 

146 for app_config in django_apps.get_app_configs(): 

147 if isinstance(app_config, PluginConfig): 

148 try: 

149 _context.update(app_config.get_jinja_context()) 

150 except Exception: 

151 logger.exception("Plugin %r raised an exception in get_jinja_context()", app_config.name) 

152 

153 if context is not None: 

154 _context.update(context) 

155 

156 return _context 

157 

158 def clean(self): 

159 super().clean() 

160 

161 params = self.environment_params or {} 

162 for key, value in params.items(): 

163 # finalize is deprecated: block new use but preserve existing stored values 

164 if key == 'finalize': 

165 raise ValidationError({ 

166 'environment_params': _( 

167 'The "{key}" parameter is deprecated and may not be set on new or modified templates.' 

168 ).format(key=key) 

169 }) 

170 if key not in JINJA_ENV_PARAMS_ALLOWED: 

171 raise ValidationError({ 

172 'environment_params': _( 

173 '"{key}" is not a permitted Jinja2 environment parameter.' 

174 ).format(key=key) 

175 }) 

176 allowed = JINJA_ENV_PARAMS_ALLOWED[key] 

177 if type(allowed) is dict: 

178 if value not in allowed: 

179 raise ValidationError({ 

180 'environment_params': _( 

181 'Invalid value "{value}" for parameter "{key}". ' 

182 'Allowed values are: {allowed}' 

183 ).format( 

184 value=value, 

185 key=key, 

186 allowed=', '.join(sorted(allowed.keys())) 

187 ) 

188 }) 

189 

190 @staticmethod 

191 def _filter_environment_params(params): 

192 """ 

193 Return a copy of params with only permitted keys. Keys not in the allowlist are 

194 stripped, except 'finalize' which is a deprecated legacy carve-out. 

195 """ 

196 return { 

197 key: value for key, value in params.items() 

198 if key in JINJA_ENV_PARAMS_ALLOWED or key == 'finalize' 

199 } 

200 

201 @staticmethod 

202 def _resolve_mapped_params(params): 

203 """ 

204 Resolve allowlisted params that have a value mapping (e.g. undefined class names) 

205 to their Python objects via direct dict lookup. Returns a new dict with resolved values; 

206 unresolved params are passed through unchanged. 

207 """ 

208 resolved = {} 

209 for name, value in params.items(): 

210 allowed = JINJA_ENV_PARAMS_ALLOWED.get(name) 

211 if type(allowed) is dict and value in allowed: 

212 resolved[name] = allowed[value] 

213 else: 

214 resolved[name] = value 

215 return resolved 

216 

217 @staticmethod 

218 def _resolve_finalize(params): 

219 """ 

220 Legacy carve-out: resolve the deprecated 'finalize' parameter via import_string(). 

221 Existing templates with finalize continue to work; new use is blocked by clean(). 

222 """ 

223 if 'finalize' in params and type(params['finalize']) is str: 

224 return {**params, 'finalize': import_string(params['finalize'])} 

225 return params 

226 

227 def get_environment_params(self): 

228 """ 

229 Pre-processing of any defined Jinja environment parameters. 

230 """ 

231 # Shallow-copy so resolved values don't replace the strings on the model field. 

232 params = dict(self.environment_params or {}) 

233 params = self._filter_environment_params(params) 

234 params = self._resolve_mapped_params(params) 

235 params = self._resolve_finalize(params) 

236 return params 

237 

238 def render(self, context=None, queryset=None): 

239 """ 

240 Render the template with the provided context. The context is passed to the Jinja2 environment as a dictionary. 

241 """ 

242 context = self.get_context(context=context, queryset=queryset) 

243 env_params = self.get_environment_params() 

244 debug = getattr(self, 'debug', False) 

245 output = render_jinja2(self.template_code, context, env_params, getattr(self, 'data_file', None), debug=debug) 

246 

247 # Replace CRLF-style line terminators 

248 output = output.replace('\r\n', '\n') 

249 

250 return output 

251 

252 def render_to_response(self, context=None, queryset=None): 

253 output = self.render(context=context, queryset=queryset) 

254 mime_type = self.mime_type or DEFAULT_MIME_TYPE 

255 

256 # Build the response 

257 response = HttpResponse(output, content_type=mime_type) 

258 

259 if self.as_attachment: 

260 extension = f'.{self.file_extension}' if self.file_extension else '' 

261 if self.file_name: 

262 filename = self.file_name 

263 elif queryset is not None: 

264 filename = filename_from_model(queryset.model) 

265 elif context: 

266 filename = filename_from_object(context) 

267 else: 

268 filename = "output" 

269 filename = f'{filename}{extension}' 

270 response['Content-Disposition'] = content_disposition_header(as_attachment=True, filename=filename) 

271 

272 return response