# Generated by Django 5.0.11 on 2025-01-27 12:58

import uuid
import pickle  # nosec
from django.core import signing
from django.contrib.auth import BACKEND_SESSION_KEY, HASH_SESSION_KEY, SESSION_KEY
from django.db import migrations, models
import django.db.models.deletion
from django.conf import settings
from django.contrib.sessions.backends.cache import KEY_PREFIX
from django.utils.timezone import now, timedelta
from authentik.lib.migrations import progress_bar
from authentik.root.middleware import ClientIPMiddleware


SESSION_CACHE_ALIAS = "default"


class PickleSerializer:
    """
    Simple wrapper around pickle to be used in signing.dumps()/loads() and
    cache backends.
    """

    def __init__(self, protocol=None):
        self.protocol = pickle.HIGHEST_PROTOCOL if protocol is None else protocol

    def dumps(self, obj):
        """Pickle data to be stored in redis"""
        return pickle.dumps(obj, self.protocol)

    def loads(self, data):
        """Unpickle data to be loaded from redis"""
        try:
            return pickle.loads(data)  # nosec
        except Exception:
            return {}


def _migrate_session(
    apps,
    db_alias,
    session_key,
    session_data,
    expires,
):
    Session = apps.get_model("authentik_core", "Session")
    OldAuthenticatedSession = apps.get_model("authentik_core", "OldAuthenticatedSession")
    AuthenticatedSession = apps.get_model("authentik_core", "AuthenticatedSession")

    old_auth_session = (
        OldAuthenticatedSession.objects.using(db_alias).filter(session_key=session_key).first()
    )

    args = {
        "session_key": session_key,
        "expires": expires,
        "last_ip": ClientIPMiddleware.default_ip,
        "last_user_agent": "",
        "session_data": {},
    }
    for k, v in session_data.items():
        if k == "authentik/stages/user_login/last_ip":
            args["last_ip"] = v
        elif k in ["last_user_agent", "last_used"]:
            args[k] = v
        elif args in [SESSION_KEY, BACKEND_SESSION_KEY, HASH_SESSION_KEY]:
            pass
        else:
            args["session_data"][k] = v
    if old_auth_session:
        args["last_user_agent"] = old_auth_session.last_user_agent
        args["last_used"] = old_auth_session.last_used

    args["session_data"] = pickle.dumps(args["session_data"])
    session = Session.objects.using(db_alias).create(**args)

    if old_auth_session:
        AuthenticatedSession.objects.using(db_alias).create(
            session=session,
            user=old_auth_session.user,
            uuid=old_auth_session.uuid,
        )


def migrate_redis_sessions(apps, schema_editor):
    from django.core.cache import caches

    db_alias = schema_editor.connection.alias
    cache = caches[SESSION_CACHE_ALIAS]

    # Not a redis cache, skipping
    if not hasattr(cache, "keys"):
        return

    print("\nMigrating Redis sessions to database, this might take a couple of minutes...")
    for key, session_data in progress_bar(cache.get_many(cache.keys(f"{KEY_PREFIX}*")).items()):
        _migrate_session(
            apps=apps,
            db_alias=db_alias,
            session_key=key.removeprefix(KEY_PREFIX),
            session_data=session_data,
            expires=now() + timedelta(seconds=cache.ttl(key)),
        )


def migrate_database_sessions(apps, schema_editor):
    DjangoSession = apps.get_model("sessions", "Session")
    db_alias = schema_editor.connection.alias

    print("\nMigration database sessions, this might take a couple of minutes...")
    for django_session in progress_bar(DjangoSession.objects.using(db_alias).all()):
        session_data = signing.loads(
            django_session.session_data,
            salt="django.contrib.sessions.SessionStore",
            serializer=PickleSerializer,
        )
        _migrate_session(
            apps=apps,
            db_alias=db_alias,
            session_key=django_session.session_key,
            session_data=session_data,
            expires=django_session.expire_date,
        )


class Migration(migrations.Migration):

    dependencies = [
        ("sessions", "0001_initial"),
        ("authentik_core", "0045_rename_new_identifier_usersourceconnection_identifier_and_more"),
        ("authentik_providers_oauth2", "0027_accesstoken_authentik_p_expires_9f24a5_idx_and_more"),
        ("authentik_providers_rac", "0006_connectiontoken_authentik_p_expires_91f148_idx_and_more"),
    ]

    operations = [
        # Rename AuthenticatedSession to OldAuthenticatedSession
        migrations.RenameModel(
            old_name="AuthenticatedSession",
            new_name="OldAuthenticatedSession",
        ),
        migrations.RenameIndex(
            model_name="oldauthenticatedsession",
            new_name="authentik_c_expires_cf4f72_idx",
            old_name="authentik_c_expires_08251d_idx",
        ),
        migrations.RenameIndex(
            model_name="oldauthenticatedsession",
            new_name="authentik_c_expirin_c1f17f_idx",
            old_name="authentik_c_expirin_9cd839_idx",
        ),
        migrations.RenameIndex(
            model_name="oldauthenticatedsession",
            new_name="authentik_c_expirin_e04f5d_idx",
            old_name="authentik_c_expirin_195a84_idx",
        ),
        migrations.RenameIndex(
            model_name="oldauthenticatedsession",
            new_name="authentik_c_session_a44819_idx",
            old_name="authentik_c_session_d0f005_idx",
        ),
        migrations.RunSQL(
            sql="ALTER INDEX authentik_core_authenticatedsession_user_id_5055b6cf RENAME TO authentik_core_oldauthenticatedsession_user_id_5055b6cf",
            reverse_sql="ALTER INDEX authentik_core_oldauthenticatedsession_user_id_5055b6cf RENAME TO authentik_core_authenticatedsession_user_id_5055b6cf",
        ),
        # Create new Session and AuthenticatedSession models
        migrations.CreateModel(
            name="Session",
            fields=[
                (
                    "session_key",
                    models.CharField(
                        max_length=40, primary_key=True, serialize=False, verbose_name="session key"
                    ),
                ),
                ("expires", models.DateTimeField(default=None, null=True)),
                ("expiring", models.BooleanField(default=True)),
                ("session_data", models.BinaryField(verbose_name="session data")),
                ("last_ip", models.GenericIPAddressField()),
                ("last_user_agent", models.TextField(blank=True)),
                ("last_used", models.DateTimeField(auto_now=True)),
            ],
            options={
                "default_permissions": [],
                "verbose_name": "Session",
                "verbose_name_plural": "Sessions",
            },
        ),
        migrations.AddIndex(
            model_name="session",
            index=models.Index(fields=["expires"], name="authentik_c_expires_d2f607_idx"),
        ),
        migrations.AddIndex(
            model_name="session",
            index=models.Index(fields=["expiring"], name="authentik_c_expirin_7c2cfb_idx"),
        ),
        migrations.AddIndex(
            model_name="session",
            index=models.Index(
                fields=["expiring", "expires"], name="authentik_c_expirin_1ab2e4_idx"
            ),
        ),
        migrations.AddIndex(
            model_name="session",
            index=models.Index(
                fields=["expires", "session_key"], name="authentik_c_expires_c49143_idx"
            ),
        ),
        migrations.CreateModel(
            name="AuthenticatedSession",
            fields=[
                (
                    "session",
                    models.OneToOneField(
                        on_delete=django.db.models.deletion.CASCADE,
                        primary_key=True,
                        serialize=False,
                        to="authentik_core.session",
                    ),
                ),
                ("uuid", models.UUIDField(default=uuid.uuid4, unique=True)),
                (
                    "user",
                    models.ForeignKey(
                        on_delete=django.db.models.deletion.CASCADE, to=settings.AUTH_USER_MODEL
                    ),
                ),
            ],
            options={
                "verbose_name": "Authenticated Session",
                "verbose_name_plural": "Authenticated Sessions",
            },
        ),
        migrations.RunPython(
            code=migrate_redis_sessions,
            reverse_code=migrations.RunPython.noop,
        ),
        migrations.RunPython(
            code=migrate_database_sessions,
            reverse_code=migrations.RunPython.noop,
        ),
    ]
