From 61d500946e4d5bdf27420091d7176e999365d568 Mon Sep 17 00:00:00 2001 From: Dustin Pianalto Date: Tue, 22 Jun 2021 01:29:38 -0800 Subject: [PATCH] Add Auto Role Feature --- cmd/goff/main.go | 3 +- internal/exts/roles/roles.go | 133 +++++++++++++++++++ internal/exts/user_management/member_join.go | 16 +++ internal/postgres/postfixes.go | 21 +++ internal/postgres/roles.go | 86 ++++++++++++ internal/services/services.go | 4 +- role.go | 21 +++ 7 files changed, 282 insertions(+), 2 deletions(-) create mode 100644 internal/postgres/roles.go create mode 100644 role.go diff --git a/cmd/goff/main.go b/cmd/goff/main.go index 7881321..a4f82a5 100644 --- a/cmd/goff/main.go +++ b/cmd/goff/main.go @@ -51,6 +51,7 @@ func main() { us := &postgres.UserService{DB: postgres.DB} gs := &postgres.GuildService{DB: postgres.DB} + rs := &postgres.RoleService{DB: postgres.DB} //prefixes := []string{ // "Go.", @@ -71,7 +72,7 @@ func main() { // Add Command Handlers exts.AddCommandHandlers(&manager) - services.InitalizeServices(us, gs) + services.InitalizeServices(us, gs, rs) //if _, ok := handler.Commands["help"]; !ok { // handler.AddDefaultHelpCommand() diff --git a/internal/exts/roles/roles.go b/internal/exts/roles/roles.go index e7f0a01..9c4323f 100644 --- a/internal/exts/roles/roles.go +++ b/internal/exts/roles/roles.go @@ -7,7 +7,9 @@ import ( "github.com/bwmarrin/discordgo" "github.com/dustinpianalto/disgoman" + "github.com/dustinpianalto/goff" "github.com/dustinpianalto/goff/internal/postgres" + "github.com/dustinpianalto/goff/internal/services" ) var MakeRoleSelfAssignableCommand = &disgoman.Command{ @@ -254,3 +256,134 @@ func memberHasRole(m *discordgo.Member, id string) bool { } return false } + +var MakeAutoRoleCommand = &disgoman.Command{ + Name: "make-auto-role", + Aliases: []string{"makear"}, + Description: "Marks the passed in role to auto add to new members.", + OwnerOnly: false, + Hidden: false, + RequiredPermissions: disgoman.PermissionManageServer, + Invoke: makeAutoRoleCommandFunc, +} + +func makeAutoRoleCommandFunc(ctx disgoman.Context, args []string) { + var roleString string + var roleID string + if len(args) > 0 { + roleString = strings.Join(args, " ") + if strings.HasPrefix(roleString, "<@&") && strings.HasSuffix(roleString, ">") { + roleID = roleString[3 : len(roleString)-1] + } else if _, err := strconv.Atoi(roleString); err == nil { + roleID = roleString + } else { + for _, role := range ctx.Guild.Roles { + if roleString == role.Name { + roleID = role.ID + } + } + } + } + fmt.Println(roleID) + var r *discordgo.Role + var err error + if r, err = ctx.Session.State.Role(ctx.Guild.ID, roleID); err != nil { + ctx.CommandManager.ErrorChannel <- disgoman.CommandError{ + Context: ctx, + Message: "Can't find that Role.", + Error: err, + } + return + } + role, err := services.RoleService.Role(r.ID) + if err != nil { + role = &goff.Role{ + ID: r.ID, + IsModerator: false, + IsAdmin: false, + SelfAssignable: false, + AutoRole: true, + Guild: ctx.Guild.ID, + } + role, err = services.RoleService.AddRole(role) + if err != nil { + ctx.CommandManager.ErrorChannel <- disgoman.CommandError{ + Context: ctx, + Message: "Error making an Auto Role", + Error: err, + } + return + } + } else { + err = services.RoleService.MakeAutoRole(role) + if err != nil { + ctx.CommandManager.ErrorChannel <- disgoman.CommandError{ + Context: ctx, + Message: "Error making an Auto Role", + Error: err, + } + return + } + } + +} + +var RemoveAutoRoleCommand = &disgoman.Command{ + Name: "remove-auto-role", + Aliases: []string{"removear"}, + Description: "Remove role from those added to new members.", + OwnerOnly: false, + Hidden: false, + RequiredPermissions: disgoman.PermissionManageServer, + Invoke: removeAutoRoleCommandFunc, +} + +func removeAutoRoleCommandFunc(ctx disgoman.Context, args []string) { + var roleString string + var roleID string + if len(args) > 0 { + roleString = strings.Join(args, " ") + if strings.HasPrefix(roleString, "<@&") && strings.HasSuffix(roleString, ">") { + roleID = roleString[3 : len(roleString)-1] + } else if _, err := strconv.Atoi(roleString); err == nil { + roleID = roleString + } else { + for _, role := range ctx.Guild.Roles { + if roleString == role.Name { + roleID = role.ID + } + } + } + } + fmt.Println(roleID) + var r *discordgo.Role + var err error + if r, err = ctx.Session.State.Role(ctx.Guild.ID, roleID); err != nil { + ctx.CommandManager.ErrorChannel <- disgoman.CommandError{ + Context: ctx, + Message: "Can't find that Role.", + Error: err, + } + return + } + role, err := services.RoleService.Role(r.ID) + if err != nil { + ctx.CommandManager.ErrorChannel <- disgoman.CommandError{ + Context: ctx, + Message: "Role was not configured as an Auto Role", + Error: err, + } + return + } else { + err = services.RoleService.RemoveAutoRole(role) + if err != nil { + ctx.CommandManager.ErrorChannel <- disgoman.CommandError{ + Context: ctx, + Message: "Error removing Auto Role", + Error: err, + } + return + } + } + +} diff --git a/internal/exts/user_management/member_join.go b/internal/exts/user_management/member_join.go index b50d28f..f8ba4b9 100644 --- a/internal/exts/user_management/member_join.go +++ b/internal/exts/user_management/member_join.go @@ -41,6 +41,22 @@ func OnGuildMemberAdd(s *discordgo.Session, member *discordgo.GuildMemberAdd) { if err != nil { log.Println("Error adding user to guild: ", err) } + roles, err := services.RoleService.GetAutoRoles(member.GuildID) + if err != nil { + log.Println("Error getting Auto Join Roles: ", err) + } + for _, r := range roles { + role, err := s.State.Role(member.GuildID, r.ID) + if err != nil { + log.Println("Error getting role: ", err) + continue + } + err = s.GuildMemberRoleAdd(member.GuildID, member.User.ID, role.ID) + if err != nil { + log.Println("Error adding Role to member: ", err) + continue + } + } } func OnGuildMemberRemove(s *discordgo.Session, member *discordgo.GuildMemberRemove) { diff --git a/internal/postgres/postfixes.go b/internal/postgres/postfixes.go index 07224d9..4ae2177 100644 --- a/internal/postgres/postfixes.go +++ b/internal/postgres/postfixes.go @@ -32,6 +32,10 @@ var postfixes = []postfix{ Name: "4_Add_Multi_Column_Unique_XUsersGuilds", Invoke: addMultiColumnUniqueXUsersGuilds, }, + postfix{ + Name: "5_Add_Auto_Role_To_Roles_Table", + Invoke: addAutoRoleToRoles, + }, } func RunPostfixes() { @@ -162,3 +166,20 @@ func addMultiColumnUniqueXUsersGuilds(revert bool) error { } return nil } + +func addAutoRoleToRoles(revert bool) error { + var queryString string + if !revert { + queryString = `ALTER TABLE roles + ADD COLUMN auto_role bool DEFAULT false` + } else { + queryString = `ALTER TABLE roles + DROP COLUMN auto_role` + } + _, err := DB.Exec(queryString) + if err != nil { + log.Println(err) + return err + } + return nil +} diff --git a/internal/postgres/roles.go b/internal/postgres/roles.go new file mode 100644 index 0000000..49bfea8 --- /dev/null +++ b/internal/postgres/roles.go @@ -0,0 +1,86 @@ +package postgres + +import ( + "database/sql" + "log" + + "github.com/dustinpianalto/goff" +) + +type RoleService struct { + DB *sql.DB +} + +func (s *RoleService) Role(id string) (*goff.Role, error) { + var r goff.Role + queryString := `SELECT id, moderator, admin, self_assignable, guild_id FROM roles WHERE id = $1` + row := s.DB.QueryRow(queryString, id) + if err := row.Scan(&r.ID, &r.IsModerator, &r.IsAdmin, &r.SelfAssignable, &r.Guild); err != nil { + return nil, err + } + return &r, nil +} + +func (s *RoleService) AddRole(r *goff.Role) (*goff.Role, error) { + queryString := `INSERT INTO roles (id, moderator, admin, self_assignable, auto_role, guild) + VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT DO NOTHING` + _, err := s.DB.Exec(queryString, r.ID, r.IsModerator, r.IsAdmin, r.SelfAssignable, r.AutoRole, r.Guild) + if err != nil { + return nil, err + } + return r, nil +} + +func (s *RoleService) DeleteRole(r *goff.Role) error { + queryString := `DELETE FROM roles WHERE id = $1` + _, err := s.DB.Exec(queryString, r.ID) + return err +} + +func (s *RoleService) MakeSelfAssignable(r *goff.Role) error { + queryString := `UPDATE roles SET self_assignable = true WHERE id = $1` + _, err := s.DB.Exec(queryString, r.ID) + return err +} + +func (s *RoleService) RemoveSelfAssignable(r *goff.Role) error { + queryString := `UPDATE roles SET self_assignable = false WHERE id = $1` + _, err := s.DB.Exec(queryString, r.ID) + return err +} + +func (s *RoleService) MakeAutoRole(r *goff.Role) error { + queryString := `UPDATE roles SET auto_role = true WHERE id = $1` + _, err := s.DB.Exec(queryString, r.ID) + return err +} + +func (s *RoleService) RemoveAutoRole(r *goff.Role) error { + queryString := `UPDATE roles SET auto_role = false WHERE id = $1` + _, err := s.DB.Exec(queryString, r.ID) + return err +} + +func (s *RoleService) GetAutoRoles(gID string) ([]*goff.Role, error) { + var roles []*goff.Role + queryString := `SELECT id FROM roles WHERE guild_id = $1` + rows, err := s.DB.Query(queryString, gID) + if err != nil { + return nil, err + } + for rows.Next() { + var id string + err := rows.Scan(&id) + if err != nil { + log.Println(err) + continue + } + role, err := s.Role(id) + if err != nil { + log.Println(err) + continue + } + roles = append(roles, role) + } + return roles, nil +} diff --git a/internal/services/services.go b/internal/services/services.go index d498bd5..1ac004d 100644 --- a/internal/services/services.go +++ b/internal/services/services.go @@ -4,8 +4,10 @@ import "github.com/dustinpianalto/goff" var UserService goff.UserService var GuildService goff.GuildService +var RoleService goff.RoleService -func InitalizeServices(us goff.UserService, gs goff.GuildService) { +func InitalizeServices(us goff.UserService, gs goff.GuildService, rs goff.RoleService) { UserService = us GuildService = gs + RoleService = rs } diff --git a/role.go b/role.go new file mode 100644 index 0000000..cb16e31 --- /dev/null +++ b/role.go @@ -0,0 +1,21 @@ +package goff + +type Role struct { + ID string + IsModerator bool + IsAdmin bool + SelfAssignable bool + AutoRole bool + Guild string +} + +type RoleService interface { + Role(id string) (*Role, error) + AddRole(r *Role) (*Role, error) + DeleteRole(r *Role) error + MakeSelfAssignable(r *Role) error + RemoveSelfAssignable(r *Role) error + MakeAutoRole(r *Role) error + RemoveAutoRole(r *Role) error + GetAutoRoles(gID string) ([]*Role, error) +}