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

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 @property 

22 def list_connection_names(self): 

23 return [cn.name for cn in self.connection_names.all()] 

24 

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

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

27 

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

29 

30 if self.svg: 

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

32 soup = BeautifulSoup(f, 'xml') 

33 connection_names = [] 

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

35 for e in elements: 

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

37 connection_name, _ = ThreatModelConnectionName.objects.get_or_create(name=dc) 

38 connection_names.append(connection_name) 

39 self.connection_names.set(connection_names)