CatStation/Content.Shared/Whitelist/EntityWhitelistSystem.cs

using System.Diagnostics.CodeAnalysis;
using Content.Shared.Item;
using Content.Shared.Tag;

namespace Content.Shared.Whitelist;

public sealed partial class EntityWhitelistSystem : EntitySystem
{
    [Dependency] private TagSystem _tag = default!;

    [Dependency] private EntityQuery<ItemComponent> _itemQuery = default!;

    private CompName _itemComponentName = default;
    private CompName _tagComponentName = default;

    public override void Initialize()
    {
        base.Initialize();

        // caching for minor performance improvement
        _itemComponentName = CompName.Get<ItemComponent>(Factory);
        _tagComponentName = CompName.Get<TagComponent>(Factory);
    }

    /// <summary>
    /// Checks whether a given entity satisfies a whitelist.
    /// Returns false if the entity is null.
    /// </summary>
    public bool IsValid(EntityWhitelist list, [NotNullWhen(true)] EntityUid? uid)
    {
        return uid != null && IsValid(list, uid.Value);
    }

    /// <summary>
    /// Checks whether a given entity satisfies a whitelist.
    /// </summary>
    public bool IsValid(EntityWhitelist list, EntityUid uid)
    {
        list.Registrations ??= StringsToRegs(list.Components);

        if (list.Registrations != null)
        {
            foreach (var reg in list.Registrations)
            {
                if (EntityManager.HasComponent(uid, reg))
                {
                    if (!list.RequireAll)
                        return true;
                }
                else if (list.RequireAll)
                    return false;
            }
        }

        if (list.Sizes != null && _itemQuery.TryComp(uid, out var itemComp))
        {
            if (list.Sizes.Contains(itemComp.Size))
                return true;
        }

        if (list.Tags != null)
        {
            return list.RequireAll
                ? _tag.HasAllTags(uid, list.Tags)
                : _tag.HasAnyTag(uid, list.Tags);
        }

        return list.RequireAll;
    }

    /// The following are a list of "helper functions" that are basically the same as each other
    /// to help make code that uses EntityWhitelist a bit more readable because at the moment
    /// it is quite clunky having to write out component.Whitelist == null ? true : _whitelist.IsValid(component.Whitelist, uid)
    /// several times in a row and makes comparisons easier to read

    /// <summary>
    /// Checks whether a given entity is allowed by a whitelist and not blocked by a blacklist.
    /// If a blacklist is provided and it matches then this returns false.
    /// If a whitelist is provided and it does not match then this returns false.
    /// If either list is null it does not get checked.
    /// </summary>
    public bool CheckBoth([NotNullWhen(true)] EntityUid? uid, EntityWhitelist? blacklist = null, EntityWhitelist? whitelist = null)
    {
        if (uid == null)
            return false;

        if (blacklist != null && IsValid(blacklist, uid))
            return false;

        return whitelist == null || IsValid(whitelist, uid);
    }

    /// <summary>
    /// Helper function to determine if a whitelist is not null and the entity is on list.
    /// </summary>
    public bool IsWhitelistPass(EntityWhitelist? whitelist, EntityUid uid)
    {
        if (whitelist == null)
            return false;

        return IsValid(whitelist, uid);
    }

    /// <summary>
    /// Helper function to determine if a whitelist is not null and the entity is not on the list.
    /// </summary>
    public bool IsWhitelistFail(EntityWhitelist? whitelist, EntityUid uid)
    {
        if (whitelist == null)
            return false;

        return !IsValid(whitelist, uid);
    }

    /// <summary>
    /// Helper function to determine if a whitelist is either null or the entity is on the list.
    /// </summary>
    public bool IsWhitelistPassOrNull(EntityWhitelist? whitelist, EntityUid uid)
    {
        if (whitelist == null)
            return true;

        return IsValid(whitelist, uid);
    }

    /// <summary>
    /// Helper function to determine if a whitelist is either null or the entity is not on the list.
    /// </summary>
    public bool IsWhitelistFailOrNull(EntityWhitelist? whitelist, EntityUid uid)
    {
        if (whitelist == null)
            return true;

        return !IsValid(whitelist, uid);
    }

    private List<ComponentRegistration>? StringsToRegs(string[]? input)
    {
        if (input == null || input.Length == 0)
            return null;

        var list = new List<ComponentRegistration>(input.Length);
        foreach (var name in input)
        {
            if (Factory.TryGetRegistration(name, out var registration))
                list.Add(registration);
            else if (Factory.GetComponentAvailability(name) != ComponentAvailability.Ignore)
                Log.Error($"{nameof(StringsToRegs)} failed: Unknown component name {name} passed to EntityWhitelist!");
        }

        return list;
    }
}