From c06d4533a2a40dc3272dedb13b2bdcc78c4c5f87 Mon Sep 17 00:00:00 2001
From: Panapol Poungdech <boom.rst28@gmail.com>
Date: Tue, 5 Jul 2022 01:02:34 +0700
Subject: [PATCH] [add] module counter

---
 gate.go                               | 31 +++++++++++++
 gate_test.go                          | 66 +++++++++++++++++++++++++++
 testdata/rbac_with_domains_policy.csv |  3 ++
 3 files changed, 100 insertions(+)

diff --git a/gate.go b/gate.go
index 588b51c..2ee4a56 100644
--- a/gate.go
+++ b/gate.go
@@ -107,6 +107,23 @@ func (g Gate) AssignRoleToUser(domain, role, user string) error {
 	return err
 }
 
+func (g Gate) CountModule(domain string) map[string]int {
+	roles := g.GetRoles(domain)
+	counter := make(map[string]int)
+	moduleRole := make(map[string][]string)
+	for _, role := range roles {
+		modules := g.GetModuleRelatedByRole(domain, role)
+		moduleRole[role] = modules
+	}
+	users := g.GetAllUsersRole(domain)
+	for _, user := range users {
+		for _, module := range moduleRole[user.Role] {
+			counter[module]++
+		}
+	}
+	return counter
+}
+
 func (g Gate) RevokeRoleToUser(domain, role, user string) error {
 	_, err := g.E.RemoveGroupingPolicy(user, role, domain)
 	return err
@@ -119,3 +136,17 @@ func (g Gate) Load() error {
 func (g Gate) Save() error {
 	return g.E.SavePolicy()
 }
+
+func (g Gate) GetModuleRelatedByRole(domain, role string) []string {
+	var modules []string
+	mapModules := map[string]struct{}{}
+	policies := g.E.GetFilteredPolicy(0, role, domain, "", "")
+	for _, policy := range policies {
+		moduleName := policy[2]
+		if _, ok := mapModules[moduleName]; !ok {
+			mapModules[moduleName] = struct{}{}
+			modules = append(modules, moduleName)
+		}
+	}
+	return modules
+}
diff --git a/gate_test.go b/gate_test.go
index 941f9f4..fe8db66 100644
--- a/gate_test.go
+++ b/gate_test.go
@@ -29,6 +29,7 @@ func TestGate_GetAllUsersRole(t *testing.T) {
 		{"alice", "admin"},
 		{"bob", "reader"},
 		{"foo", "visitor"},
+		{"chalet", "visitor"},
 	}
 	assert.Equal(t, expectedD1, u1)
 
@@ -37,6 +38,7 @@ func TestGate_GetAllUsersRole(t *testing.T) {
 		{"alice", "reader2"},
 		{"bob", "admin"},
 		{"foo", "visitor2"},
+		{"chalet", "visitor2"},
 	}
 	assert.Equal(t, expectedD2, u2)
 
@@ -44,6 +46,7 @@ func TestGate_GetAllUsersRole(t *testing.T) {
 	expectedD3 := []UserRole{
 		{"alice", "visitor3"},
 		{"bob", "reader3"},
+		{"chalet", "reader3"},
 		{"foo", "admin"},
 	}
 	assert.Equal(t, expectedD3, u3)
@@ -268,3 +271,66 @@ func TestGate_AddPolicy(t *testing.T) {
 	err = g.Load()
 	assert.NoError(t, err)
 }
+
+func TestGate_CountModule(t *testing.T) {
+	a := fileadapter.NewAdapter("./testdata/rbac_with_domains_policy.csv")
+	g, err := NewGate(a)
+	assert.NoError(t, err)
+
+	c := g.CountModule("domain1")
+	assert.NoError(t, err)
+	expect := map[string]int{
+		"data2": 1,
+		"data3": 2,
+	}
+	assert.Equal(t, expect, c)
+
+	c = g.CountModule("domain2")
+	assert.NoError(t, err)
+	expect = map[string]int{
+		"data2": 1,
+	}
+	assert.Equal(t, expect, c)
+
+	c = g.CountModule("domain3")
+	assert.NoError(t, err)
+	expect = map[string]int{
+		"data2": 2,
+		"data3": 1,
+	}
+	assert.Equal(t, expect, c)
+
+	c = g.CountModule("domain5")
+	assert.NoError(t, err)
+	expect = map[string]int{}
+	assert.Equal(t, expect, c)
+}
+
+func TestGate_GetModuleRelatedByRole(t *testing.T) {
+	a := fileadapter.NewAdapter("./testdata/rbac_with_domains_policy.csv")
+	g, err := NewGate(a)
+	assert.NoError(t, err)
+
+	m := g.GetModuleRelatedByRole("domain1", "reader")
+	assert.NoError(t, err)
+	expect := []string{"data2"}
+	assert.Equal(t, expect, m)
+
+	m = g.GetModuleRelatedByRole("domain1", "visitor")
+	assert.NoError(t, err)
+	expect = []string{"data3"}
+	assert.Equal(t, expect, m)
+
+	m = g.GetModuleRelatedByRole("domain1", "writer")
+	assert.NoError(t, err)
+	expect = []string{"data1"}
+	assert.Equal(t, expect, m)
+
+	m = g.GetModuleRelatedByRole("domain1", "writer2")
+	assert.NoError(t, err)
+	assert.Nil(t, m)
+
+	m = g.GetModuleRelatedByRole("domain5", "writer2")
+	assert.NoError(t, err)
+	assert.Nil(t, m)
+}
diff --git a/testdata/rbac_with_domains_policy.csv b/testdata/rbac_with_domains_policy.csv
index ccbdf9f..5af299c 100644
--- a/testdata/rbac_with_domains_policy.csv
+++ b/testdata/rbac_with_domains_policy.csv
@@ -15,9 +15,12 @@ p, Observer3, domain3, data3, view
 g, alice, admin, domain1
 g, bob, reader, domain1
 g, foo, visitor, domain1
+g, chalet, visitor, domain1
 g, alice, reader2, domain2
 g, bob, admin, domain2
 g, foo, visitor2, domain2
+g, chalet, visitor2, domain2
 g, alice, visitor3, domain3
 g, bob, reader3, domain3
+g, chalet, reader3, domain3
 g, foo, admin, domain3
\ No newline at end of file