Mini Shell

import time
from typing import Generator

import blinker
from peewee import (
    JOIN,
    CharField,
    Check,
    DoesNotExist,
    ForeignKeyField,
    IntegerField,
)
from playhouse.shortcuts import model_to_dict

from defence360agent.model import instance, Model
from defence360agent.model.simplification import apply_order_by
from im360.internals import geo


class Country(Model):
    """Contains a single record per country, with its code and name."""

    id = CharField(primary_key=True, null=False)
    #: Country code, e.g. "US".
    code = CharField(max_length=2, unique=True, null=False)
    #: Country name, e.g. "United States".
    name = CharField(null=False)

    class Meta:
        db_table = "country"
        database = instance.db

    @classmethod
    def update_from(cls, countries):
        """
        Load country info from iterable of dicts to database
        :param countries:
        :return:
        """
        for country in countries:
            cls.insert(country).on_conflict_replace().execute()


class CountryList(Model):
    """List of Countries in user WHITE/BLACK list."""

    #: Listname for blocked countries.
    BLACK = "BLACK"
    #: Listname for whitelisted countries (not officially supported).
    WHITE = "WHITE"

    #: The link to country code and name in :class:`Country`.
    country = ForeignKeyField(Country, primary_key=True, null=False)
    #: In which list all IPs from this country should be:
    #: :attr:`BLACK` or :attr:`WHITE`.
    listname = CharField(
        null=False, constraints=[Check("listname in ('WHITE','BLACK')")]
    )
    #: Timestamp when the record was added.
    ctime = IntegerField(null=True, default=lambda: int(time.time()))
    #: Comment set by admin.
    comment = CharField(null=True)

    class Meta:
        db_table = "country_list"
        database = instance.db

    class Signals:
        added = blinker.Signal()
        deleted = blinker.Signal()

    @classmethod
    def create(cls, **kwargs):
        obj = super().create(**kwargs)
        cls.Signals.added.send(obj.listname, country_id=obj.country_id)
        return obj

    @classmethod
    def fetch_as_union(cls, **filter_args):
        return cls._fetch_filter(skip_order_by=True, **filter_args)

    @classmethod
    def _fetch_filter(
        cls,
        by_country_code=None,
        by_comment=None,
        by_ip=None,
        by_list=None,
        skip_order_by=False,
    ):
        country_code_from_ip = None
        if by_ip:
            with geo.reader() as geo_reader:
                # country_code_from_ip will be None if here partial IP is provided  # noqa: E501
                country_code_from_ip = geo_reader.get_code(by_ip)

        q = (
            CountryList.select(CountryList)
            .distinct()
            .join(Country, JOIN.INNER, on=(CountryList.country == Country.id))
        )
        if by_list:
            q = q.where(CountryList.listname.in_(by_list))

        if not skip_order_by:
            q = q.order_by(Country.code)

        # filter by optional args
        if by_country_code:
            q = q.where(Country.code == by_country_code)
        if by_comment:
            q = q.where(cls.comment.contains(by_comment))
        if by_ip:
            q = q.where(Country.code == country_code_from_ip)
        return q

    @classmethod
    def fetch_count(cls, **filter_args):
        return cls._fetch_filter(**filter_args).count()

    @classmethod
    def fetch(cls, offset=None, limit=None, order_by=None, **filter_args):
        q = cls._fetch_filter(**filter_args)
        if offset is not None:
            q = q.offset(offset)
        if limit is not None:
            q = q.limit(limit)
        if order_by is not None:
            q = apply_order_by(order_by, cls, q)
        return [model_to_dict(row) for row in q]

    @classmethod
    def get_listname(cls, country):
        try:
            return cls.get(
                country=Country.select().where(Country.code == country)
            ).listname
        except DoesNotExist:
            return None

    @classmethod
    def delete_country(cls, country: Country, listname):
        deleted = (
            CountryList.delete()
            .where(
                (CountryList.country == country)
                & (CountryList.listname == listname)
            )
            .execute()
        )
        if deleted:
            cls.Signals.deleted.send(listname, country_id=country.id)
        return deleted

    @classmethod
    def country_codes(cls, listname) -> Generator:
        """Returns generator of listed country codes."""
        q = (
            CountryList.select(Country.code)
            .distinct()
            .join(Country, JOIN.INNER, on=(CountryList.country == Country.id))
            .where(CountryList.listname == listname)
            .order_by(Country.code)
        )

        return (code for (code,) in q.tuples())