remote_billing: Sort out remote_billing_identities typing.

This does two important things:
1. Fix return type of get_identity_dict_from_session to correctly be
   Optional[Union[RemoteBillingIdentityDict, LegacyServerIdentityDict]].
   RemoteBillingIdentityDict is the type in the 8.0+ auth flow,
   LegacyServerIdentityDict is the type in old servers flow, where only
   the server uuid info is available.
2. The uuid key used in request.session["remote_billing_identities"]
   should be explicitly namespaced depending on which flow and type
   we're
   dealing with - to avoid confusion in case of collisions between a
   realm and server that have the same UUID. Such a situation should not
   occur naturally and I haven't come up with any actual exploitation
   ideas that could utilize this by manipulating your server/realm
   uuids, but it's much easier to just not think about such collision
   security implications by making them impossible.

corporate/lib/remote_billing_util.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Optional, TypedDict
+from typing import Optional, TypedDict, Union, cast
 
 from django.http import HttpRequest
 from django.utils.translation import gettext as _
@@ -28,25 +28,34 @@ class LegacyServerIdentityDict(TypedDict):
 
 def get_identity_dict_from_session(
     request: HttpRequest,
+    *,
     realm_uuid: Optional[str],
     server_uuid: Optional[str],
-) -> Optional[RemoteBillingIdentityDict]:
-    authed_uuid = realm_uuid or server_uuid
-    assert authed_uuid is not None
+) -> Optional[Union[RemoteBillingIdentityDict, LegacyServerIdentityDict]]:
+    if not (realm_uuid or server_uuid):
+        return None
 
     identity_dicts = request.session.get("remote_billing_identities")
-    if identity_dicts is not None:
-        return identity_dicts.get(authed_uuid)
+    if identity_dicts is None:
+        return None
 
-    return None
+    if realm_uuid is not None:
+        return identity_dicts.get(f"remote_realm:{realm_uuid}")
+    else:
+        assert server_uuid is not None
+        return identity_dicts.get(f"remote_server:{server_uuid}")
 
 
 def get_remote_realm_from_session(
     request: HttpRequest,
     realm_uuid: Optional[str],
-    server_uuid: Optional[str] = None,
 ) -> RemoteRealm:
-    identity_dict = get_identity_dict_from_session(request, realm_uuid, server_uuid)
+    # Cannot use isinstance with TypeDicts, to make mypy know
+    # which of the TypedDicts in the Union this is - so just cast it.
+    identity_dict = cast(
+        Optional[RemoteBillingIdentityDict],
+        get_identity_dict_from_session(request, realm_uuid=realm_uuid, server_uuid=None),
+    )
 
     if identity_dict is None:
         raise JsonableError(_("User not authenticated"))
@@ -77,7 +86,10 @@ def get_remote_server_from_session(
     request: HttpRequest,
     server_uuid: str,
 ) -> RemoteZulipServer:
-    identity_dict = get_identity_dict_from_session(request, None, server_uuid)
+    identity_dict: Optional[LegacyServerIdentityDict] = get_identity_dict_from_session(
+        request, realm_uuid=None, server_uuid=server_uuid
+    )
+
     if identity_dict is None:
         raise JsonableError(_("User not authenticated"))