# Authentik expression policy: sync a member's Discord ROLES into Authentik GROUPS at every login.
# Paste into: Customization -> Policies -> Create -> Expression Policy  (name: discord-roles-to-groups)
# Bind to BOTH flows (after discord-guild-member): default-source-enrollment and default-source-authentication.
# Requires the "guilds.members.read" scope on the Discord source (source -> Protocol settings -> Scopes:
#   guilds guilds.members.read ). Existing users re-consent once at Discord.
# The Authentik groups must already exist (Directory -> Groups). Role IDs: Discord Developer Mode ->
#   Server Settings -> Roles -> right-click a role -> Copy Role ID.
# Manual memberships are safe: the policy only removes groups it granted itself (user attribute discord_role_groups).

GUILD_ID = "1387203044084551690"                         # <- your Discord server ID
ROLE_TO_GROUP = {                                       # <- Discord role id : Authentik group name
    "1387239222070612110": "bot-users",
    "1515901670943490129": "bot-users",
#    "222222222222222222": "streamers",
#    "333333333333333333": "editors",
#    "444444444444444444": "admins",
    "555555555555555555": "git-users",
}

from authentik.core.models import Group

source = context.get("source")
if not source or getattr(source, "provider_type", "") != "discord":
    return True                                          # not a Discord login (e.g. akadmin): nothing to sync

connection = context.get("goauthentik.io/sources/connection")
if not connection:
    return True

resp = requests.get(
    f"https://discord.com/api/users/@me/guilds/{GUILD_ID}/member",
    headers={"Authorization": f"Bearer {connection.access_token}"},
    timeout=10,
)
if resp.status_code != 200:
    ak_message(f"Could not read your Discord roles ({resp.status_code}); groups unchanged.")
    return True                                          # don't block login on a Discord hiccup

roles = set(resp.json().get("roles", []))
wanted = {g for r, g in ROLE_TO_GROUP.items() if r in roles}
managed = set(ROLE_TO_GROUP.values())

# Authentication flow: the user exists -> edit memberships directly.
# Only memberships THIS policy granted earlier are ever removed (tracked in user.attributes["discord_role_groups"]),
# so a group someone was added to by hand in Authentik survives role changes and re-logins.
user = context.get("pending_user") or (request.user if request.user.is_authenticated else None)
if user and user.pk:
    attrs = user.attributes or {}
    granted = set(attrs.get("discord_role_groups", []))
    for group in Group.objects.filter(name__in=managed):
        if group.name in wanted:
            group.users.add(user)
            granted.add(group.name)
        elif group.name in granted:
            group.users.remove(user)          # role gone -> take back what the role had given
            granted.discard(group.name)
    attrs["discord_role_groups"] = sorted(granted)
    user.attributes = attrs
    user.save()
    return True

# Enrollment flow: the user is created later by the user-write stage; hand it the groups to attach and
# record them as role-granted so a later role removal can revoke exactly these.
plan = context.get("flow_plan")
if plan is not None:
    plan.context["groups"] = list(Group.objects.filter(name__in=wanted))
    prompt = plan.context.setdefault("prompt_data", {})
    prompt.setdefault("attributes", {})["discord_role_groups"] = sorted(wanted)
return True