from django.db import models

# Create your models here.

class MalwarePrediction(models.Model):
    MODEL_TYPE_CHOICES = [
        (1, 'KNeighborsClassifier'),
        (2, 'RandomForestClassifier'),
        (3, 'XGBClassifier'),
        (4, 'SGDClassifier'),
    ]

    process_name = models.CharField(max_length=255)
    process_class = models.CharField(max_length=50)
    probability_of_malware = models.FloatField()
    predicted_malware = models.CharField(max_length=50)
    model_type = models.IntegerField(choices=MODEL_TYPE_CHOICES, default=1)  # New field to store model type

    def __str__(self):
        return f"{self.process_name} - {self.predicted_malware} - {self.get_model_type_display()}"