from rest_framework import permissions
from django.contrib.auth import get_user_model

class IsOwnerOrAdmin(permissions.BasePermission):
    """
    Object-level permission to only allow owner or admin to edit an object.
    Assumes the model instance has an `owner` attribute.
    """

    def has_permission(self, request, view):
        if request.user.is_staff:
            return True

        try:
            target_user = get_user_model().objects.get(
                    username=view.kwargs['user_pk'])
            return target_user == request.user
        except:
            return False

    def has_object_permission(self, request, view, obj):
        return (obj.owner == request.user) or request.user.is_staff