From 477cab3945654e663cdb98b306172de0c605ae24 Mon Sep 17 00:00:00 2001
From: Teifion <teifion@teifion.co.uk>
Date: Sun, 25 Aug 2024 17:36:12 +0100
Subject: [PATCH] Added helper functions related to user restrictions

---
 lib/teiserver/account/libs/user_lib.ex | 43 ++++++++++++++++++++++++++
 test/account/user_lib_test.exs         | 24 ++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/lib/teiserver/account/libs/user_lib.ex b/lib/teiserver/account/libs/user_lib.ex
index 277e54f0..045746e5 100644
--- a/lib/teiserver/account/libs/user_lib.ex
+++ b/lib/teiserver/account/libs/user_lib.ex
@@ -192,6 +192,49 @@ defmodule Teiserver.Account.UserLib do
     |> maybe_decache_user()
   end
 
+  @doc """
+  Removes one or more restrictions from a user.
+
+  ## Examples
+
+      iex> unrestrict_user(user_or_user_id, ["r1", "r2"])
+      {:ok, %User{}}
+
+  """
+  @spec unrestrict_user(User.t() | User.id(), [String.t()] | String.t()) :: {:ok, User.t()} | {:error, Ecto.Changeset.t()}
+  def unrestrict_user(user_or_user_id, restrictions) when is_binary(user_or_user_id),
+    do: unrestrict_user(get_user_by_id(user_or_user_id), restrictions)
+
+  def unrestrict_user(%User{} = user, restrictions_to_remove) do
+    restrictions_to_remove = List.wrap(restrictions_to_remove)
+
+    new_restrictions = user.restrictions
+      |> Enum.filter(fn existing_restriction ->
+        not Enum.member?(restrictions_to_remove, existing_restriction)
+      end)
+
+    update_user(user, %{restrictions: new_restrictions})
+  end
+
+  @doc """
+  Adds one or more restrictions to a user
+
+  ## Examples
+
+      iex> remove_restrictions(user_or_user_id, ["r1", "r2"])
+      {:ok, %User{}}
+
+  """
+  @spec restrict_user(User.t() | User.id(), [String.t()] | String.t()) :: {:ok, User.t()} | {:error, Ecto.Changeset.t()}
+  def restrict_user(user_or_user_id, restrictions) when is_binary(user_or_user_id),
+    do: restrict_user(get_user_by_id(user_or_user_id), restrictions)
+
+  def restrict_user(%User{} = user, restrictions) do
+    new_restrictions = Enum.uniq(user.restrictions ++ List.wrap(restrictions))
+    update_user(user, %{restrictions: new_restrictions})
+  end
+
+
   # Clears the cache for a user after a successful database option
   @spec maybe_decache_user(any()) :: any()
   defp maybe_decache_user({:ok, user}) do
diff --git a/test/account/user_lib_test.exs b/test/account/user_lib_test.exs
index 98c3795a..ce059f77 100644
--- a/test/account/user_lib_test.exs
+++ b/test/account/user_lib_test.exs
@@ -106,6 +106,30 @@ defmodule Teiserver.UserLibTest do
       refute Account.valid_password?(user, "bad_password")
     end
 
+    test "restrict_user/2 and unrestrict_user/2" do
+      user = AccountFixtures.user_fixture()
+
+      # As a string
+      Account.UserLib.restrict_user(user, "OneRestriction")
+      user = Account.get_user!(user.id)
+      assert user.restrictions == ["OneRestriction"]
+
+      # As a list
+      Account.UserLib.restrict_user(user, ["TwoRestriction"])
+      user = Account.get_user!(user.id)
+      assert user.restrictions == ["OneRestriction", "TwoRestriction"]
+
+      # Two at once, one is an overlap
+      Account.UserLib.restrict_user(user, ["TwoRestriction", "ThreeRestriction"])
+      user = Account.get_user!(user.id)
+      assert user.restrictions == ["OneRestriction", "TwoRestriction", "ThreeRestriction"]
+
+      # Now remove
+      Account.UserLib.unrestrict_user(user, ["TwoRestriction", "ThreeRestriction"])
+      user = Account.get_user!(user.id)
+      assert user.restrictions == ["OneRestriction"]
+    end
+
     test "allow?/2" do
       # User must have all of the required permissions
       user = AccountFixtures.user_fixture(%{groups: ["perm1", "perm2"]})