Add bool value to check what type of alert the program will send

vigilate · Jul 22, 2016 · 4326b68 · 4326b68
1 parent 57ec153
commit 4326b68
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/vigilate_backend/models.py b/vigilate_backend/models.py
@@ -86,8 +86,11 @@ class UserPrograms(models.Model):
     poste = models.ForeignKey('Station')
     cpe = models.ForeignKey('vulnerability_manager.Cpe')
     sms_score = models.IntegerField(validators=[MinValueValidator(0), MaxValueValidator(10)], default=0)
+    sms_enabled = models.BooleanField(default=True)
     email_score = models.IntegerField(validators=[MinValueValidator(0), MaxValueValidator(10)], default=0)
+    email_enabled = models.BooleanField(default=True)
     web_score = models.IntegerField(validators=[MinValueValidator(0), MaxValueValidator(10)], default=0)
+    web_enabled = models.BooleanField(default=True)
     alert_type_default = models.BooleanField(default=True)
 
     def is_vulnerable(self):

diff --git a/vigilate_backend/serializers.py b/vigilate_backend/serializers.py
@@ -41,7 +41,7 @@ class UserProgramsSerializer(serializers.ModelSerializer):
 
     class Meta:
         model = models.UserPrograms
-        fields = ('id', 'program_name', 'program_version', 'minimum_score', 'user_id', 'poste', 'cpe', 'sms_score', 'email_score', 'web_score', 'alert_type_default', 'is_vulnerable')
+        fields = ('id', 'program_name', 'program_version', 'minimum_score', 'user_id', 'poste', 'cpe', 'sms_score', 'email_score', 'web_score', 'sms_enabled', 'email_enabled', 'web_enabled', 'alert_type_default', 'is_vulnerable')
 
     def create(self, validated_data):
         """Create an user program
@@ -60,6 +60,9 @@ def update(self, instance, validated_data):
         instance.sms_score = validated_data.get('sms_score', instance.sms_score)
         instance.email_score = validated_data.get('email_score', instance.email_score)
         instance.web_score = validated_data.get('web_score', instance.web_score)
+        instance.sms_enabled = validated_data.get('sms_enabled', instance.sms_enabled)
+        instance.email_enabled = validated_data.get('email_enabled', instance.email_enabled)
+        instance.web_enabled = validated_data.get('web_enabled', instance.web_enabled) 
         instance.alert_type_default = validated_data.get('alert_type_default', instance.alert_type_default)
         instance.save()
         return instance