33 lines
1.3 KiB
Python
33 lines
1.3 KiB
Python
from django.db import models
|
|
from Device .models import Devices
|
|
from Accounts.models import UserProfile
|
|
|
|
# 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()}"
|
|
|
|
|
|
class MalwarePredictionsDevice(models.Model):
|
|
device = models.ForeignKey(Devices, on_delete=models.CASCADE,null=True)
|
|
user = models.ForeignKey(UserProfile, on_delete=models.CASCADE , null=True) # Add this field to reference the user
|
|
file_path = models.FileField(upload_to='malware_predictions/')
|
|
uploaded_at = models.DateTimeField(auto_now_add=True)
|
|
|
|
def __str__(self):
|
|
return f"Prediction for {self.device.device_name} by {self.user.user.username} at {self.uploaded_at}"
|
|
|