devicehub-django/device/models.py

from django.db import models, connection

from utils.constants import STR_SM_SIZE, STR_SIZE, STR_EXTEND_SIZE, ALGOS
from evidence.models import Annotation, Evidence
from user.models import User
from lot.models import DeviceLot


class Device:
    class Types(models.TextChoices):
        DESKTOP = "Desktop"
        LAPTOP = "Laptop"
        SERVER = "Server"
        GRAPHICCARD = "GraphicCard"
        HARDDRIVE = "HardDrive"
        SOLIDSTATEDRIVE = "SolidStateDrive"
        MOTHERBOARD = "Motherboard"
        NETWORKADAPTER = "NetworkAdapter"
        PROCESSOR = "Processor"
        RAMMODULE = "RamModule"
        SOUNDCARD = "SoundCard"
        DISPLAY = "Display"
        BATTERY = "Battery"
        CAMERA = "Camera"

    def __init__(self, *args, **kwargs):
        # the id is the chid of the device
        self.id = kwargs["id"]
        self.pk = self.id
        self.algorithm = None
        self.owner = None
        self.annotations =  []
        self.hids = []
        self.uuids = []
        self.evidences = []
        self.lots = []
        self.last_evidence = None
        self.get_last_evidence()

    def initial(self):
        self.get_annotations()
        self.get_uuids()
        self.get_hids()
        self.get_evidences()
        self.get_lots()

    def get_annotations(self):
        if self.annotations:
            return self.annotations

        self.annotations = Annotation.objects.filter(
            type=Annotation.Type.SYSTEM,
            value=self.id
        ).order_by("-created")

        if self.annotations.count():
            self.algorithm = self.annotations[0].key
            self.owner = self.annotations[0].owner

        return self.annotations

    def get_user_annotations(self):
        if not self.uuids:
            self.get_uuids()

        annotations = Annotation.objects.filter(
            uuid__in=self.uuids,
            owner=self.owner,
            type=Annotation.Type.USER
        )
        return annotations

    def get_user_documents(self):
        if not self.uuids:
            self.get_uuids()

        annotations = Annotation.objects.filter(
            uuid__in=self.uuids,
            owner=self.owner,
            type=Annotation.Type.DOCUMENT
        )
        return annotations

    def get_uuids(self):
        for a in self.get_annotations():
            if a.uuid not in self.uuids:
                self.uuids.append(a.uuid)

    def get_hids(self):
        annotations = self.get_annotations()

        self.hids = annotations.filter(
            type=Annotation.Type.SYSTEM,
            key__in=ALGOS.keys(),
        ).values_list("value", flat=True)

    def get_evidences(self):
        if not self.uuids:
            self.get_uuids()

        self.evidences = [Evidence(u) for u in self.uuids]

    def get_last_evidence(self):
        annotations = self.get_annotations()
        if annotations:
            annotation = annotations.first()
        self.last_evidence = Evidence(annotation.uuid)

    def last_uuid(self):
        return self.uuids[0]

    def get_lots(self):
        self.lots = [x.lot for x in DeviceLot.objects.filter(device_id=self.id)]

    @classmethod
    def get_unassigned(cls, user, offset=0, limit=None):

        sql = """
              SELECT DISTINCT t1.value from evidence_annotation as t1
                left join lot_devicelot as t2 on t1.value = t2.device_id
              where t2.device_id is null and owner_id=={user} and type=={type}
        """.format(
            user=user.id,
            type=Annotation.Type.SYSTEM,
        )
        if limit:
            sql += " limit {} offset {}".format(int(limit), int(offset))

        sql += ";"

        annotations = []
        with connection.cursor() as cursor:
            cursor.execute(sql)
            annotations = cursor.fetchall()

        devices = [cls(id=x[0]) for x in annotations]
        count = cls.get_unassigned_count(user)
        return devices, count


    @classmethod
    def get_unassigned_count(cls, user):

        sql = """
              SELECT count(DISTINCT t1.value) from evidence_annotation as t1
                left join lot_devicelot as t2 on t1.value = t2.device_id
              where t2.device_id is null and owner_id=={user} and type=={type};
        """.format(
            user=user.id,
            type=Annotation.Type.SYSTEM,
        )

        with connection.cursor() as cursor:
            cursor.execute(sql)
            return cursor.fetchall()[0][0]

    @property
    def is_websnapshot(self):
        if not self.last_evidence:
            self.get_last_evidence()
        return self.last_evidence.doc['type'] == "WebSnapshot" 
    
    @property
    def last_user_evidence(self):
        if not self.last_evidence:
            self.get_last_evidence()
        return self.last_evidence.doc['kv'].items()
    
    @property
    def manufacturer(self):
        if not self.last_evidence:
            self.get_last_evidence()
        return self.last_evidence.doc['device']['manufacturer']

    @property
    def type(self):
        if not self.last_evidence:
            self.get_last_evidence()
        return self.last_evidence.doc['device']['type']

    @property
    def model(self):
        if not self.last_evidence:
            self.get_last_evidence()
        return self.last_evidence.doc['device']['model']

    @property
    def type(self):
        if not self.last_evidence:
            self.get_last_evidence()
        return self.last_evidence.doc['device']['type']