Coverage for product_risk_suite/threat_model/models.py: 98%

1from django.db import models 

2from django.utils.text import slugify 

3 

4from bs4 import BeautifulSoup 

5 

6from product.models import Product 

7from .shared_models import ThreatModelConnectionName 

8 

9class ThreatModel(models.Model): 

10 id = models.AutoField(primary_key=True) 

11 product = models.ForeignKey(Product, on_delete=models.PROTECT) 

12 slug = models.SlugField(default="", null=False, unique=True) 

13 title = models.CharField(max_length=200) 

14 description = models.TextField(blank=True, null=True) 

15 svg = models.FileField(upload_to='uploads/') 

16 connection_names = models.ManyToManyField(ThreatModelConnectionName, blank=True) 

17 

18 def __str__(self): 

19 return self.title 

20 

21 def __remove_svg_attribute(self, svg, tag, attr): 

22 for element in svg.findAll(tag): 

23 if attr in element.attrs: 

24 del element.attrs[attr] 

25 return svg 

26 

27 def __remove_dark_mode_from_style(self, svg): 

28 for element in svg.findAll('svg'): 

29 if 'style' in element.attrs: 

30 s = element.attrs['style'] 

31 element.attrs['style'] = s.replace('color-scheme: light dark;', '') 

32 return svg 

33 

34 def get_clean_svg_content(self): 

35 svg_content = None 

36 with self.svg.open("r") as f: 

37 soup = BeautifulSoup(f, 'xml') 

38 soup = self.__remove_svg_attribute(soup, 'svg', 'content') 

39 svg_content = str(self.__remove_dark_mode_from_style(soup)) 

40 return svg_content 

41 

42 @property 

43 def list_connection_names(self): 

44 return [cn.tech_name for cn in self.connection_names.all()] 

45 

46 def save(self, *args, **kwargs): 

47 self.slug = slugify(f"{self.product.slug} {self.title}") 

48 

49 super().save(*args, **kwargs) 

50 

51 if self.svg: 

52 with self.svg.open('r') as f: 

53 soup = BeautifulSoup(f, 'xml') 

54 connection_names = [] 

55 elements = soup.find_all(attrs={"data-cell-id": True}) 

56 for e in elements: 

57 dc = e['data-cell-id'] 

58 connection_name, _ = ThreatModelConnectionName.objects.get_or_create(product_id=self.product.id, tech_name=dc) 

59 connection_names.append(connection_name) 

60 self.connection_names.set(connection_names)