allianceauth/allianceauth/eveonline/autogroups/models.py

import logging
from typing import ClassVar
from django.db import models, transaction
from django.contrib.auth.models import Group, User
from django.core.exceptions import ObjectDoesNotExist

from allianceauth.authentication.models import State
from allianceauth.eveonline.models import EveCorporationInfo, EveAllianceInfo

logger = logging.getLogger(__name__)


def get_users_for_state(state: State):
    return User.objects.select_related('profile').prefetch_related('profile__main_character')\
            .filter(profile__state_id=state.pk)


class AutogroupsConfigManager(models.Manager):
    def update_groups_for_state(self, state: State):
        """
        Update all the Group memberships for the users
        who have State
        :param state: State to update for
        :return:
        """
        users = get_users_for_state(state)
        for config in self.filter(states=state):
            logger.debug("in state loop")
            for user in users:
                logger.debug(f"in user loop for {user}")
                config.update_group_membership_for_user(user)

    def update_groups_for_user(self, user: User, state: State = None):
        """
        Update the Group memberships for the given users state
        :param user: User to update for
        :param state: State to update user for
        :return:
        """
        if state is None:
            state = user.profile.state
        for config in self.filter(states=state):
            # grant user new groups for their state
            config.update_group_membership_for_user(user)
        for config in self.exclude(states=state):
            # ensure user does not have groups from previous state
            config.remove_user_from_alliance_groups(user)
            config.remove_user_from_corp_groups(user)


class AutogroupsConfig(models.Model):
    OPT_TICKER = 'ticker'
    OPT_NAME = 'name'
    NAME_OPTIONS = (
        (OPT_TICKER, 'Ticker'),
        (OPT_NAME, 'Full name'),
    )

    states = models.ManyToManyField(State, related_name='autogroups')

    corp_groups = models.BooleanField(default=False, help_text="Setting this to false will delete all the created groups.")
    corp_group_prefix = models.CharField(max_length=50, default='Corp ', blank=True)
    corp_name_source = models.CharField(max_length=20, choices=NAME_OPTIONS, default=OPT_NAME)

    alliance_groups = models.BooleanField(default=False, help_text="Setting this to false will delete all the created groups.")
    alliance_group_prefix = models.CharField(max_length=50, default='Alliance ', blank=True)
    alliance_name_source = models.CharField(max_length=20, choices=NAME_OPTIONS, default=OPT_NAME)

    corp_managed_groups = models.ManyToManyField(
        Group, through='ManagedCorpGroup', related_name='corp_managed_config',
        help_text='A list of corporation groups created and maintained by this AutogroupConfig. You should not edit this list unless you know what you\'re doing.')

    alliance_managed_groups = models.ManyToManyField(
        Group, through='ManagedAllianceGroup', related_name='alliance_managed_config',
        help_text='A list of alliance groups created and maintained by this AutogroupConfig. You should not edit this list unless you know what you\'re doing.')

    replace_spaces = models.BooleanField(default=False)
    replace_spaces_with = models.CharField(
        max_length=10, default='', blank=True,
        help_text='Any spaces in the group name will be replaced with this.')

    objects: ClassVar[AutogroupsConfigManager] = AutogroupsConfigManager()

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __repr__(self):
        return self.__class__.__name__

    def __str__(self):
        return 'States: ' + (' '.join(list(self.states.all().values_list('name', flat=True))) if self.pk else str(None))

    def update_all_states_group_membership(self):
        list(map(self.update_group_membership_for_state, self.states.all()))

    def update_group_membership_for_state(self, state: State):
        list(map(self.update_group_membership_for_user, get_users_for_state(state)))

    @transaction.atomic
    def update_group_membership_for_user(self, user: User):
        self.update_alliance_group_membership(user)
        self.update_corp_group_membership(user)

    def user_entitled_to_groups(self, user: User) -> bool:
        try:
            return user.profile.state in self.states.all()
        except ObjectDoesNotExist:
            return False

    @transaction.atomic
    def update_alliance_group_membership(self, user: User):
        group = None
        try:
            if not self.alliance_groups or not self.user_entitled_to_groups(user):
                logger.debug(f'User {user} does not have required state for alliance group membership')
                return
            else:
                alliance = user.profile.main_character.alliance
                if alliance is None:
                    logger.debug(f'User {user} alliance is None, cannot update group membership')
                    return
                group = self.get_alliance_group(alliance)
        except EveAllianceInfo.DoesNotExist:
            logger.debug(f'User {user} main characters alliance does not exist in the database. Creating.')
            alliance = EveAllianceInfo.objects.create_alliance(user.profile.main_character.alliance_id)
            group = self.get_alliance_group(alliance)
        except AttributeError:
            logger.warning(f'User {user} does not have a main character. Group membership not updated')
        finally:
            self.remove_user_from_alliance_groups(user, except_group=group)
            if group is not None:
                logger.debug(f'Adding user {user} to alliance group {group}')
                user.groups.add(group)

    @transaction.atomic
    def update_corp_group_membership(self, user: User):
        group = None
        try:
            if not self.corp_groups or not self.user_entitled_to_groups(user):
                logger.debug(f'User {user} does not have required state for corp group membership')
            else:
                corp = user.profile.main_character.corporation
                group = self.get_corp_group(corp)
        except EveCorporationInfo.DoesNotExist:
            logger.debug(f'User {user} main characters corporation does not exist in the database. Creating.')
            corp = EveCorporationInfo.objects.create_corporation(user.profile.main_character.corporation_id)
            group = self.get_corp_group(corp)
        except AttributeError:
            logger.warning(f'User {user} does not have a main character. Group membership not updated')
        finally:
            self.remove_user_from_corp_groups(user, except_group=group)
            if group is not None:
                logger.debug(f'Adding user {user} to corp group {group}')
                user.groups.add(group)

    @transaction.atomic
    def remove_user_from_alliance_groups(self, user: User, except_group: Group = None):
        remove_groups = user.groups.filter(pk__in=self.alliance_managed_groups.all().values_list('pk', flat=True))
        if except_group is not None:
            remove_groups = remove_groups.exclude(pk=except_group.pk)
        list(map(user.groups.remove, remove_groups))

    @transaction.atomic
    def remove_user_from_corp_groups(self, user: User, except_group: Group = None):
        remove_groups = user.groups.filter(pk__in=self.corp_managed_groups.all().values_list('pk', flat=True))
        if except_group is not None:
            remove_groups = remove_groups.exclude(pk=except_group.pk)
        list(map(user.groups.remove, remove_groups))

    def get_alliance_group(self, alliance: EveAllianceInfo) -> Group:
        return self.create_alliance_group(alliance)

    def get_corp_group(self, corp: EveCorporationInfo) -> Group:
        return self.create_corp_group(corp)

    @transaction.atomic
    def create_alliance_group(self, alliance: EveAllianceInfo) -> Group:
        group, created = Group.objects.get_or_create(name=self.get_alliance_group_name(alliance))
        ManagedAllianceGroup.objects.get_or_create(group=group, config=self, alliance=alliance)
        return group

    @transaction.atomic
    def create_corp_group(self, corp: EveCorporationInfo) -> Group:
        group, created = Group.objects.get_or_create(name=self.get_corp_group_name(corp))
        ManagedCorpGroup.objects.get_or_create(group=group, config=self, corp=corp)
        return group

    def delete_alliance_managed_groups(self):
        """
        Deletes ALL managed alliance groups
        """
        for g in self.alliance_managed_groups.all():
            g.delete()

    def delete_corp_managed_groups(self):
        """
        Deletes ALL managed corp groups
        """
        for g in self.corp_managed_groups.all():
            g.delete()

    def get_alliance_group_name(self, alliance: EveAllianceInfo) -> str:
        if self.alliance_name_source == self.OPT_TICKER:
            name = alliance.alliance_ticker
        elif self.alliance_name_source == self.OPT_NAME:
            name = alliance.alliance_name
        else:
            raise NameSourceException('Not a valid name source')
        return self._replace_spaces(self.alliance_group_prefix + name)

    def get_corp_group_name(self, corp: EveCorporationInfo) -> str:
        if self.corp_name_source == self.OPT_TICKER:
            name = corp.corporation_ticker
        elif self.corp_name_source == self.OPT_NAME:
            name = corp.corporation_name
        else:
            raise NameSourceException('Not a valid name source')
        return self._replace_spaces(self.corp_group_prefix + name)

    def _replace_spaces(self, name: str) -> str:
        """
        Replace the spaces in the given name based on the config
        :param name: name to replace spaces in
        :return: name with spaces replaced with the configured character(s) or unchanged if configured
        """
        if self.replace_spaces:
            return name.strip().replace(' ', str(self.replace_spaces_with))
        return name


class ManagedGroup(models.Model):
    group = models.ForeignKey(Group, on_delete=models.CASCADE)
    config = models.ForeignKey(AutogroupsConfig, on_delete=models.CASCADE)

    class Meta:
        abstract = True

    def __str__(self):
        return "Managed Group: %s" % self.group.name

class ManagedCorpGroup(ManagedGroup):
    corp = models.ForeignKey(EveCorporationInfo, on_delete=models.CASCADE)


class ManagedAllianceGroup(ManagedGroup):
    alliance = models.ForeignKey(EveAllianceInfo, on_delete=models.CASCADE)


class NameSourceException(Exception):
    pass