Review comments addressed

- Removed loop to find add/remove class

Jira: ASCYB-3003

Signed-off-by: Mahesh Thumar <maheshth@cybage.com>