Source.Python forums

from events import Event
from players.entity import Player
from weapons.entity import Weapon
from entities.helpers import index_from_pointer
from entities.hooks import EntityCondition, EntityPreHook
from filters.players import PlayerIter
from memory import make_object
from messages import SayText2
from colors import GREEN, LIGHT_GREEN

limit = 2


class T_RestrictPlayer(Player):
	caching = True

	def __init__(self, index):
		super().__init__(index)
		self.is_limit_reached = 0
		self.is_picked_weapon = False


class CT_RestrictPlayer(Player):
	caching = True

	def __init__(self, index):
		super().__init__(index)
		self.is_limit_reached = 0
		self.is_picked_weapon = False

@Event('player_death')
def player_death(args):
	player = Player.from_userid(args['userid'])
	userid = player.userid
	if player.team == 3: # Is a counter-terrorist
		ct = CT_RestrictPlayer.from_userid(userid)
		if ct.is_picked_weapon: # Have the user picked up the weapon
			for ct_player in PlayerIter('ct'):
				ct_pla = CT_RestrictPlayer.from_userid(userid)
				if ct_pla.is_limit_reached >= 0:
					ct_pla.is_limit_reached -= 1
					ct.is_picked_weapon = False
	elif player.team == 2: # Is a terrorist
		t = T_RestrictPlayer.from_userid(userid)
		if t.is_picked_weapon: # Have the user picked up the weapon
			for t_player in PlayerIter('t'):
				t_pla = T_RestrictPlayer.from_userid(userid)
				if t_pla.is_limit_reached >= 0:
					t_pla.is_limit_reached -= 1
					t.is_picked_weapon = False

@EntityPreHook(EntityCondition.is_player, 'buy_internal')
def pre_buy(args):
	player = Player(index_from_pointer(args[0]))
	weapon = args[1]
	userid = player.userid
	if weapon in ['sg550', 'g3sg1']:
		SayText2(f"{GREEN}[Weapon Restrict] -> {LIGHT_GREEN}You can't buy {GREEN}autosnipers!").send(player.index)
		return False
	elif weapon == 'awp':
		if player.team == 3: # Is counter-terrorist
			ct = CT_RestrictPlayer.from_userid(userid)
			if not ct.is_picked_weapon:
				ct.is_picked_weapon = True
				for ct_player in PlayerIter('ct'):
					ct_pla = CT_RestrictPlayer.from_userid(userid)
					if not ct_pla.is_limit_reached >= limit:
						ct_pla.is_limit_reached += 1
					else:
						SayText2(f"{GREEN}[Weapon Restrict] -> {LIGHT_GREEN}You can't buy {GREEN}awp!. {LIGHT_GREEN}The limit 2 of 2 is {GREEN}reached!").send(player.index)
						return False
		if player.team == 2: # Is terrorist
			t = T_RestrictPlayer.from_userid(userid)
			if not t.is_picked_weapon:
				t.is_picked_weapon = True
				for t_player in PlayerIter('t'):
					t_pla = T_RestrictPlayer.from_userid(userid)
					if not t_pla.is_limit_reached >= limit:
						t_pla.is_limit_reached += 1
					else:
						SayText2(f"{GREEN}[Weapon Restrict] -> {LIGHT_GREEN}You can't buy {GREEN}awp!. {LIGHT_GREEN}The limit 2 of 2 is {GREEN}reached!").send(player.index)
						return False

@EntityPreHook(EntityCondition.is_human_player, 'drop_weapon')
@EntityPreHook(EntityCondition.is_bot_player, 'drop_weapon')
def pre_weapon_drop(args):
	player = Player(index_from_pointer(args[0]))
	weapon = make_object(Weapon, args[1])
	userid = player.userid
	if weapon == 'weapon_awp':
		if player.team == 3: # Is counter-terrorist
			ct = CT_RestrictPlayer.from_userid(userid)
			if ct.is_picked_weapon:
				ct.is_picked_weapon = False
				for ct_player in PlayerIter('ct'):
					ct_pla = CT_RestrictPlayer.from_userid(userid)
					if ct_pla.is_limit_reached >= 0:
						ct_pla.is_limit_reached -= 1
						ct_pla.is_picked_weapon = False
			if player.team == 2: # Is terrorist
				t = T_RestrictPlayer.from_userid(userid)
				if t.is_picked_weapon:
					t.is_picked_weapon = False
					for t_player in PlayerIter('t'):
						t_pla = T_RestrictPlayer.from_userid(userid)
						if t_pla.is_limit_reached >= 0:
							t_pla.is_limit_reached -= 1
							t_pla.is_picked_weapon = False

@EntityPreHook(EntityCondition.is_human_player, 'bump_weapon')
@EntityPreHook(EntityCondition.is_bot_player, 'bump_weapon')
def pre_pickup(args):
	player = Player(index_from_pointer(args[0]))
	weapon = Weapon(index_from_pointer(args[1])).classname.replace('weapon_', '', 1)
	userid = player.userid
	if weapon in ['sg550', 'g3sg1']:
		SayText2(f"{GREEN}[Weapon Restrict] -> {LIGHT_GREEN}You can't pick up {GREEN}autosnipers!").send(player.index)
		return False
	elif weapon == 'awp':
		if player.team == 3: # Is counter-terrorist
			ct = CT_RestrictPlayer.from_userid(userid)
			if not ct.is_picked_weapon:
				ct.is_picked_weapon = True
				for ct_player in PlayerIter('ct'):
					ct_pla = CT_RestrictPlayer.from_userid(userid)
					if not ct_pla.is_limit_reached >= limit:
						ct_pla.is_limit_reached += 1
					else:
						SayText2(f"{GREEN}[Weapon Restrict] -> {LIGHT_GREEN}You can't pickup {GREEN}awp!. {LIGHT_GREEN}The limit 2 of 2 is {GREEN}reached!").send(player.index)
						return False
		if player.team == 2: # Is terrorist
			t = T_RestrictPlayer.from_userid(userid)
			if not t.is_picked_weapon:
				t.is_picked_weapon = True
				for t_player in PlayerIter('t'):
					t_pla = T_RestrictPlayer.from_userid(userid)
					if not t_pla.is_limit_reached >= limit:
						t_pla.is_limit_reached += 1
					else:
						SayText2(f"{GREEN}[Weapon Restrict] -> {LIGHT_GREEN}You can't pickup {GREEN}awp!. {LIGHT_GREEN}The limit 2 of 2 is {GREEN}reached!").send(player.index)
						return False

from filters.entities import EntityIter
from weapons.manager import weapon_manager
from weapons.restrictions import WeaponRestrictionHandler
from core import GAME_NAME


TEAM_RESTRICTED_WEAPONS = {
    'sg550': 0,
    'g3sg1': 0,
    'awp': 2,
}


class MyWeaponRestrictionHandler(WeaponRestrictionHandler):
    def on_player_purchasing_weapon(self, player, weapon_name):
        return self.can_player_have_weapon(player, weapon_name)

    def on_player_bumping_weapon(self, player, weapon_name):
        return self.can_player_have_weapon(player, weapon_name)

    @staticmethod
    def can_player_have_weapon(player, weapon_name):
        max_count = TEAM_RESTRICTED_WEAPONS.get(weapon_name)
        if max_count is None:
            return True

        if not max_count:
            return False

        team_index = player.team_index
        weapon_full_name = weapon_manager[weapon_name].name
        current_count = 0
        for entity in EntityIter(weapon_full_name):
            if getattr(entity.owner, 'team_index', None) == team_index:
                current_count += 1
                if current_count == max_count:
                    return False

        return True

weapon_restriction_handler = MyWeaponRestrictionHandler()