zulip/zilencer/models.py

# https://github.com/typeddjango/django-stubs/issues/1698
# mypy: disable-error-code="explicit-override"

from dataclasses import dataclass
from datetime import datetime
from typing import List, Tuple

from django.conf import settings
from django.core.exceptions import ValidationError
from django.db import models
from django.db.models import Max, Q, UniqueConstraint
from django.utils.timezone import now as timezone_now
from typing_extensions import override

from analytics.models import BaseCount
from zerver.lib.rate_limiter import RateLimitedObject
from zerver.lib.rate_limiter import rules as rate_limiter_rules
from zerver.models import AbstractPushDeviceToken, AbstractRealmAuditLog, Realm, UserProfile


def get_remote_server_by_uuid(uuid: str) -> "RemoteZulipServer":
    try:
        return RemoteZulipServer.objects.get(uuid=uuid)
    except ValidationError:
        raise RemoteZulipServer.DoesNotExist


class RemoteZulipServer(models.Model):
    """Each object corresponds to a single remote Zulip server that is
    registered for the Mobile Push Notifications Service via
    `manage.py register_server`.
    """

    UUID_LENGTH = 36
    API_KEY_LENGTH = 64
    HOSTNAME_MAX_LENGTH = 128
    VERSION_MAX_LENGTH = 128

    # The unique UUID (`zulip_org_id`) and API key (`zulip_org_key`)
    # for this remote server registration.
    uuid = models.UUIDField(unique=True)
    api_key = models.CharField(max_length=API_KEY_LENGTH)

    # The hostname and contact details are not verified/trusted. Thus,
    # they primarily exist so that we can communicate with the
    # maintainer of a server about abuse problems.
    hostname = models.CharField(max_length=HOSTNAME_MAX_LENGTH)
    contact_email = models.EmailField(blank=True, null=False)
    last_updated = models.DateTimeField("last updated", auto_now=True)
    last_version = models.CharField(max_length=VERSION_MAX_LENGTH, null=True)

    # Whether the server registration has been deactivated.
    deactivated = models.BooleanField(default=False)

    # Plan types for self-hosted customers
    PLAN_TYPE_SELF_HOSTED = 1
    PLAN_TYPE_COMMUNITY = 100
    PLAN_TYPE_BUSINESS = 101
    PLAN_TYPE_ENTERPRISE = 102

    # The current billing plan for the remote server, similar to Realm.plan_type.
    plan_type = models.PositiveSmallIntegerField(default=PLAN_TYPE_SELF_HOSTED)

    # This is not synced with the remote server, but only filled for sponsorship requests.
    org_type = models.PositiveSmallIntegerField(
        default=Realm.ORG_TYPES["unspecified"]["id"],
        choices=[(t["id"], t["name"]) for t in Realm.ORG_TYPES.values()],
    )

    @override
    def __str__(self) -> str:
        return f"{self.hostname} {str(self.uuid)[0:12]}"

    def format_requester_for_logs(self) -> str:
        return "zulip-server:" + str(self.uuid)


class RemotePushDeviceToken(AbstractPushDeviceToken):
    """Like PushDeviceToken, but for a device connected to a remote server."""

    server = models.ForeignKey(RemoteZulipServer, on_delete=models.CASCADE)
    # The user id on the remote server for this device
    user_id = models.BigIntegerField(null=True)
    user_uuid = models.UUIDField(null=True)

    remote_realm = models.ForeignKey("RemoteRealm", on_delete=models.SET_NULL, null=True)

    class Meta:
        unique_together = [
            # These indexes rely on the property that in Postgres,
            # NULL != NULL in the context of unique indexes, so multiple
            # rows with the same values in these columns can exist
            # if one of them is NULL.
            ("server", "user_id", "kind", "token"),
            ("server", "user_uuid", "kind", "token"),
        ]

    @override
    def __str__(self) -> str:
        return f"{self.server!r} {self.user_id}"


class RemoteRealm(models.Model):
    """
    Each object corresponds to a single remote Realm that is using the
    Mobile Push Notifications Service via `manage.py register_server`.
    """

    server = models.ForeignKey(RemoteZulipServer, on_delete=models.CASCADE)

    # The unique UUID and secret for this realm.
    uuid = models.UUIDField(unique=True)
    uuid_owner_secret = models.TextField()

    # Value obtained's from the remote server's realm.host.
    host = models.TextField()

    name = models.TextField(default="")

    authentication_methods = models.JSONField(default=dict)

    org_type = models.PositiveSmallIntegerField(
        default=Realm.ORG_TYPES["unspecified"]["id"],
        choices=[(t["id"], t["name"]) for t in Realm.ORG_TYPES.values()],
    )

    # The fields below are analogical to RemoteZulipServer fields.

    last_updated = models.DateTimeField("last updated", auto_now=True)

    # Whether the realm registration has been deactivated.
    registration_deactivated = models.BooleanField(default=False)
    # Whether the realm has been deactivated on the remote server.
    realm_deactivated = models.BooleanField(default=False)

    # When the realm was created on the remote server.
    realm_date_created = models.DateTimeField()

    # Plan types for self-hosted customers
    PLAN_TYPE_SELF_HOSTED = 1
    PLAN_TYPE_COMMUNITY = 100
    PLAN_TYPE_BUSINESS = 101
    PLAN_TYPE_ENTERPRISE = 102

    # The current billing plan for the remote server, similar to Realm.plan_type.
    plan_type = models.PositiveSmallIntegerField(default=PLAN_TYPE_SELF_HOSTED, db_index=True)

    @override
    def __str__(self) -> str:
        return f"{self.host} {str(self.uuid)[0:12]}"


class RemoteRealmBillingUser(models.Model):
    remote_realm = models.ForeignKey(RemoteRealm, on_delete=models.CASCADE)

    # The .uuid of the UserProfile on the remote server
    user_uuid = models.UUIDField()
    full_name = models.TextField(default="")
    email = models.EmailField()

    TOS_VERSION_BEFORE_FIRST_LOGIN = UserProfile.TOS_VERSION_BEFORE_FIRST_LOGIN
    tos_version = models.TextField(default=TOS_VERSION_BEFORE_FIRST_LOGIN)


class RemoteZulipServerAuditLog(AbstractRealmAuditLog):
    """Audit data associated with a remote Zulip server (not specific to a
    realm).  Used primarily for tracking registration and billing
    changes for self-hosted customers.

    In contrast with RemoteRealmAuditLog, which has a copy of data
    that is generated on the client Zulip server, this table is the
    authoritative storage location for the server's history.
    """

    server = models.ForeignKey(RemoteZulipServer, on_delete=models.CASCADE)

    @override
    def __str__(self) -> str:
        return f"{self.server!r} {self.event_type} {self.event_time} {self.id}"


class RemoteRealmAuditLog(AbstractRealmAuditLog):
    """Synced audit data from a remote Zulip server, used primarily for
    billing.  See RealmAuditLog and AbstractRealmAuditLog for details.
    """

    server = models.ForeignKey(RemoteZulipServer, on_delete=models.CASCADE)

    # With modern Zulip servers, we can link to the RemoteRealm object.
    remote_realm = models.ForeignKey(RemoteRealm, on_delete=models.CASCADE, null=True)
    # For pre-8.0 servers, we might only have the realm ID and thus no
    # RemoteRealm object yet. We will eventually be able to drop this
    # column once all self-hosted servers have upgraded in favor of
    # just using the foreign key everywhere.
    realm_id = models.IntegerField(null=True, blank=True)
    # The remote_id field lets us deduplicate data from the remote server
    remote_id = models.IntegerField(null=True)

    @override
    def __str__(self) -> str:
        return f"{self.server!r} {self.event_type} {self.event_time} {self.id}"

    class Meta:
        constraints = [
            models.UniqueConstraint(
                fields=["server", "remote_id"],
                name="zilencer_remoterealmauditlog_server_remote",
            ),
        ]
        indexes = [
            models.Index(
                fields=["server", "realm_id", "remote_id"],
                name="zilencer_remoterealmauditlog_server_realm_remote",
            ),
        ]


class BaseRemoteCount(BaseCount):
    server = models.ForeignKey(RemoteZulipServer, on_delete=models.CASCADE)
    # The remote_id field is the id value of the corresponding *Count object
    # on the remote server.
    # It lets us deduplicate data from the remote server.
    # Note: Some counts don't come from the remote server, but rather
    # are stats we track on the bouncer server itself, pertaining to the remote server.
    # E.g. mobile_pushes_received::day. Such counts will set this field to None.
    remote_id = models.IntegerField(null=True)

    class Meta:
        abstract = True


class RemoteInstallationCount(BaseRemoteCount):
    class Meta:
        constraints = [
            UniqueConstraint(
                fields=["server", "property", "subgroup", "end_time"],
                condition=Q(subgroup__isnull=False),
                name="unique_remote_installation_count",
            ),
            UniqueConstraint(
                fields=["server", "property", "end_time"],
                condition=Q(subgroup__isnull=True),
                name="unique_remote_installation_count_null_subgroup",
            ),
            UniqueConstraint(
                fields=["server", "remote_id"],
                # As noted above, remote_id may be null, so we only
                # enforce uniqueness if it isn't.  This is not
                # technically necessary, since null != null, but it
                # makes the property more explicit.
                condition=Q(remote_id__isnull=False),
                name="unique_remote_installation_count_server_id_remote_id",
            ),
        ]

    @override
    def __str__(self) -> str:
        return f"{self.property} {self.subgroup} {self.value}"


# We can't subclass RealmCount because we only have a realm_id here, not a foreign key.
class RemoteRealmCount(BaseRemoteCount):
    realm_id = models.IntegerField(null=True)
    # Certain RemoteRealmCount will be counts tracked on the bouncer server directly, about
    # stats pertaining to a realm on a remote server. For such objects, we will link to
    # the corresponding RemoteRealm object that the remote server registered with us.
    # In the future we may be able to link all RemoteRealmCount objects to a RemoteRealm,
    # including the RemoteRealmCount objects are results of just syncing the RealmCount
    # table from the remote server.
    remote_realm = models.ForeignKey(RemoteRealm, on_delete=models.CASCADE, null=True)

    class Meta:
        constraints = [
            UniqueConstraint(
                fields=["server", "realm_id", "property", "subgroup", "end_time"],
                condition=Q(subgroup__isnull=False),
                name="unique_remote_realm_installation_count",
            ),
            UniqueConstraint(
                fields=["server", "realm_id", "property", "end_time"],
                condition=Q(subgroup__isnull=True),
                name="unique_remote_realm_installation_count_null_subgroup",
            ),
            UniqueConstraint(
                fields=["server", "remote_id"],
                # As with RemoteInstallationCount above, remote_id may
                # be null; since null != null, this condition is not
                # strictly necessary, but serves to make the property
                # more explicit.
                condition=Q(remote_id__isnull=False),
                name="unique_remote_realm_installation_count_server_id_remote_id",
            ),
        ]
        indexes = [
            models.Index(
                fields=["property", "end_time"],
                name="zilencer_remoterealmcount_property_end_time_506a0b38_idx",
            ),
        ]

    @override
    def __str__(self) -> str:
        return f"{self.server!r} {self.realm_id} {self.property} {self.subgroup} {self.value}"


class RateLimitedRemoteZulipServer(RateLimitedObject):
    def __init__(
        self, remote_server: RemoteZulipServer, domain: str = "api_by_remote_server"
    ) -> None:
        # Remote servers can only make API requests regarding push notifications
        # which requires ZILENCER_ENABLED and of course can't happen on API endpoints
        # inside Tornado.
        assert not settings.RUNNING_INSIDE_TORNADO
        assert settings.ZILENCER_ENABLED

        self.uuid = str(remote_server.uuid)
        self.domain = domain
        super().__init__()

    @override
    def key(self) -> str:
        return f"{type(self).__name__}:<{self.uuid}>:{self.domain}"

    @override
    def rules(self) -> List[Tuple[int, int]]:
        return rate_limiter_rules[self.domain]


@dataclass
class RemoteCustomerUserCount:
    guest_user_count: int
    non_guest_user_count: int


def get_remote_server_guest_and_non_guest_count(
    server_id: int, event_time: datetime = timezone_now()
) -> RemoteCustomerUserCount:
    # For each realm hosted on the server, find the latest audit log
    # entry indicating the number of active users in that realm.
    realm_last_audit_log_ids = (
        RemoteRealmAuditLog.objects.filter(
            server_id=server_id,
            event_type__in=RemoteRealmAuditLog.SYNCED_BILLING_EVENTS,
            event_time__lte=event_time,
        )
        # Important: extra_data is empty for some pre-2020 audit logs
        # prior to the introduction of realm_user_count_by_role
        # logging. Meanwhile, modern Zulip servers using
        # bulk_create_users to create the users in the system bot
        # realm also generate such audit logs. Such audit logs should
        # never be the latest in a normal realm.
        .exclude(extra_data={})
        .values("realm_id")
        .annotate(max_id=Max("id"))
        .values_list("max_id", flat=True)
    )

    extra_data_list = RemoteRealmAuditLog.objects.filter(
        id__in=list(realm_last_audit_log_ids)
    ).values_list("extra_data", flat=True)

    # Now we add up the user counts from the different realms.
    guest_count = 0
    non_guest_count = 0
    for extra_data in extra_data_list:
        humans_count_dict = extra_data[RemoteRealmAuditLog.ROLE_COUNT][
            RemoteRealmAuditLog.ROLE_COUNT_HUMANS
        ]
        for role_type in UserProfile.ROLE_TYPES:
            if role_type == UserProfile.ROLE_GUEST:
                guest_count += humans_count_dict[str(role_type)]
            else:
                non_guest_count += humans_count_dict[str(role_type)]

    return RemoteCustomerUserCount(
        non_guest_user_count=non_guest_count, guest_user_count=guest_count
    )