from __future__ import unicode_literals

from mongoengine import *

from core.entities import Entity

from core.database import StringListField


class MalwareFamily(Document):

    name = StringField(required=True, unique=True)

    def __unicode__(self):
        return self.name


class Malware(Entity):

    aliases = ListField(StringField(), verbose_name="Aliases")
    family = ReferenceField(MalwareFamily, verbose_name="Family")

    DISPLAY_FIELDS = Entity.DISPLAY_FIELDS + [("aliases", "Aliases"),
                                              ("family", "Family")]

    @classmethod
    def get_form(klass):
        form = Entity.get_form(override=klass)
        form.aliases = StringListField("Aliases")
        return form

    def info(self):
        i = Entity.info(self)
        i['family'] = self.family.name if self.family else None
        i['type'] = "Malware"
        return i

    def generate_tags(self):
        tags = [self.name.lower()]
        if self.family is not None:
            tags.append(self.family.name.lower())
        return tags
