Merge pull request 'fix(bot): Check if user exist before perform action' (#6) from Starthur/bot:main into main

Reviewed-on: 42CTF/bot#6

Dockerfile

@@ -0,0 +1,12 @@
+FROM golang:1.19
+
+WORKDIR /usr/src/app/
+
+COPY go.mod go.sum ./
+RUN go mod download && go mod verify
+
+COPY . .
+
+RUN go build -v -o /usr/local/bin/bot ./...
+
+CMD ["bot"]

README.md

@@ -11,5 +11,22 @@ AUTH_TOKEN=<api_42CTF_token>
 
 And then run it with :
 ```
-docker-compose up -d
+docker-compose up --build -d
 ```
+By default it will run on 42CTF Guild and with RolesID from this server.
+You can run :
+```
+docker-compose run bot_discord [--wait WAIT] [--force] [--dryrun]
+```
+
+## Arguments
+|      arg     | Default | Usage                                         |
+|:------------:|:-------:|:---------------------------------------------:|
+| --guild / -g |  `nil`  | Guild ID                                      |
+|    --top1    |  `nil`  | TOP1 Role ID                                  |
+|   --top10    |  `nil`  | TOP10 Role ID                                 |
+|   --top50    |  `nil`  | TOP10 Role ID                                 |
+|   --wait     |  `1`    | Minutes to wait between runs                  |
+|   --force    | `false` | Force update of all users                     |
+|   --dryrun   | `false` | Display log but don't execute discord command |
+

api.go

@@ -0,0 +1,33 @@
+package main
+
+import (
+	"encoding/json"
+	"io"
+	"net/http"
+	"time"
+)
+
+func getUsers42CTF() (map[string]int, error) {
+	var jsonParse map[string]int
+
+	client := http.Client{ Timeout: time.Second * 2}
+	req, err := http.NewRequest(http.MethodGet, api, nil)
+	if err != nil {
+		return nil, err
+	}
+	res, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer res.Body.Close()
+	body, err := io.ReadAll(res.Body)
+	if err != nil {
+		return nil, err
+	}
+	err = json.Unmarshal([]byte(body), &jsonParse)
+	if err != nil {
+		return nil, err
+	}
+	return jsonParse, nil
+}
+

discord.go

@@ -0,0 +1,28 @@
+package main
+
+import (
+	"github.com/sirupsen/logrus"
+	"github.com/bwmarrin/discordgo"
+)
+
+func removeRole(discordClient *discordgo.Session, guildID string, userID string, roleID string, dryrun bool) {
+	if ! dryrun {
+		err := discordClient.GuildMemberRoleRemove(guild, userID, roleID)
+		if err != nil {
+			logrus.WithFields(logrus.Fields{
+				"id": userID,
+			}).Error(err.Error())
+		}
+	}
+}
+
+func addRole(discordClient *discordgo.Session, guildID string, userID string, roleID string, dryrun bool) {
+	if ! dryrun {
+		err := discordClient.GuildMemberRoleAdd(guild, userID, roleID)
+		if err != nil {
+			logrus.WithFields(logrus.Fields{
+				"id": userID,
+			}).Error(err.Error())
+		}
+	}
+}

docker-compose.yaml

@@ -2,6 +2,7 @@ version: "2"
 services:
   bot_discord:
     restart: always
-    build: ./src
+    build: .
     container_name: bot_discord
     env_file: ./env_file
+    entrypoint: bot --guild=606162827274616845 --top1=798638767359524875 --top10=801787467064672286 --top50=803729539145924649

go.mod

@@ -0,0 +1,17 @@
+module gitea.42ctf.org/42CTF/bot
+
+go 1.20
+
+require (
+	github.com/bwmarrin/discordgo v0.27.0
+	golang.org/x/exp v0.0.0-20230203172020-98cc5a0785f9
+)
+
+require (
+	github.com/alexflint/go-arg v1.4.3 // indirect
+	github.com/alexflint/go-scalar v1.1.0 // indirect
+	github.com/gorilla/websocket v1.4.2 // indirect
+	github.com/sirupsen/logrus v1.9.0 // indirect
+	golang.org/x/crypto v0.0.0-20210421170649-83a5a9bb288b // indirect
+	golang.org/x/sys v0.1.0 // indirect
+)

go.sum

@@ -0,0 +1,30 @@
+github.com/alexflint/go-arg v1.4.3 h1:9rwwEBpMXfKQKceuZfYcwuc/7YY7tWJbFsgG5cAU/uo=
+github.com/alexflint/go-arg v1.4.3/go.mod h1:3PZ/wp/8HuqRZMUUgu7I+e1qcpUbvmS258mRXkFH4IA=
+github.com/alexflint/go-scalar v1.1.0 h1:aaAouLLzI9TChcPXotr6gUhq+Scr8rl0P9P4PnltbhM=
+github.com/alexflint/go-scalar v1.1.0/go.mod h1:LoFvNMqS1CPrMVltza4LvnGKhaSpc3oyLEBUZVhhS2o=
+github.com/bwmarrin/discordgo v0.27.0 h1:4ZK9KN+rGIxZ0fdGTmgdCcliQeW8Zhu6MnlFI92nf0Q=
+github.com/bwmarrin/discordgo v0.27.0/go.mod h1:NJZpH+1AfhIcyQsPeuBKsUtYrRnjkyu0kIVMCHkZtRY=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
+github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/sirupsen/logrus v1.9.0 h1:trlNQbNUG3OdDrDil03MCb1H2o9nJ1x4/5LYw7byDE0=
+github.com/sirupsen/logrus v1.9.0/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+golang.org/x/crypto v0.0.0-20210421170649-83a5a9bb288b h1:7mWr3k41Qtv8XlltBkDkl8LoP3mpSgBW8BUoxtEdbXg=
+golang.org/x/crypto v0.0.0-20210421170649-83a5a9bb288b/go.mod h1:T9bdIzuCu7OtxOm1hfPfRQxPLYneinmdGuTeoZ9dtd4=
+golang.org/x/exp v0.0.0-20230203172020-98cc5a0785f9 h1:frX3nT9RkKybPnjyI+yvZh6ZucTZatCCEm9D47sZ2zo=
+golang.org/x/exp v0.0.0-20230203172020-98cc5a0785f9/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc=
+golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
+golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.1.0 h1:kunALQeHf1/185U1i0GOB/fy1IPRDDpuoOOqRReG57U=
+golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

main.go

@@ -0,0 +1,65 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/alexflint/go-arg"
+	"github.com/bwmarrin/discordgo"
+	"github.com/sirupsen/logrus"
+)
+
+type UserRole struct {
+	Rank        int
+	CurrentRole string
+}
+
+// Global var (I know, and i don't care)
+var api string = os.Getenv("API_URL")
+var guild string
+var rolesID map[string]string
+var Users map[string]*UserRole
+
+var args struct {
+	Guild   string `arg:"-g,required,env:GUILD_ID"`
+	Top1    string `arg:"env:TOP_1"`
+	Top10   string `arg:"env:TOP_10"`
+	Top50   string `arg:"env:TOP_50"`
+	Wait    int    `default:"2" help:"Minutes between run"`
+	Force   bool
+	DryRun  bool
+}
+
+func main() {
+	fmt.Println(len(os.Args), os.Args)
+	arg.MustParse(&args)
+	guild = args.Guild
+	rolesID = make(map[string]string)
+	rolesID["top1"] = args.Top1
+	rolesID["top10"] = args.Top10
+	rolesID["top50"] = args.Top50
+	discord, err := discordgo.New("Bot " + os.Getenv("DISCORD_TOKEN"))
+	if err != nil {
+		logrus.Fatal(err.Error())
+	}
+	if err := populateUserRole(discord); err != nil {
+		logrus.Fatal(err.Error())
+	}
+
+	if args.DryRun {
+		logrus.Warn("Running in Dry Run mode")
+	}
+
+	if args.Force {
+		logrus.Warn("Force update of all users")
+		forceUpdate(discord, args.DryRun)
+	}
+
+	for {
+		logrus.Info("Start update")
+		updateUserRole(discord, args.DryRun)
+		logrus.Info("End update, sleep for 2 minutes")
+		time.Sleep(time.Duration(args.Wait) * time.Minute)
+	}
+}

rank.go

@@ -0,0 +1,131 @@
+package main
+
+import (
+	"github.com/bwmarrin/discordgo"
+	"github.com/sirupsen/logrus"
+	"golang.org/x/exp/slices"
+)
+
+// Populate Users map for the first launch
+// This variable will be keep as sort of cache for all
+// operation
+func populateUserRole(discordClient *discordgo.Session) error {
+	Users = make(map[string]*UserRole)
+	APIResponse, err := getUsers42CTF()
+	if err != nil {
+		return err
+	}
+
+	for id, rank := range APIResponse {
+		tmp, err := discordClient.GuildMember(guild, id)
+		if err == nil {
+			if slices.Contains(tmp.Roles, rolesID["top1"]) {
+				Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top1"]}
+			} else if slices.Contains(tmp.Roles, rolesID["top10"]) {
+				Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top10"]}
+			} else if slices.Contains(tmp.Roles, rolesID["top50"]) {
+				Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top50"]}
+			} else {
+				Users[id] = &UserRole{Rank: rank, CurrentRole: ""}
+			}
+		} else {
+			logrus.WithFields(logrus.Fields{
+				"id": id,
+				"rank": rank,
+			}).Error(err.Error())
+		}
+	}
+	return nil
+}
+
+func updateUserRole(discordClient *discordgo.Session, dryrun bool) error {
+	APIResponse, err := getUsers42CTF()
+	if err != nil {
+		return err
+	}
+	for id, rank := range APIResponse {
+		tmp, err := discordClient.GuildMember(guild, id)
+		if err == nil {
+			if slices.Contains(tmp.Roles, rolesID["top1"]) {
+				Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top1"]}
+			} else if slices.Contains(tmp.Roles, rolesID["top10"]) {
+				Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top10"]}
+			} else if slices.Contains(tmp.Roles, rolesID["top50"]) {
+				Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top50"]}
+			} else {
+				Users[id] = &UserRole{Rank: rank, CurrentRole: ""}
+			}
+		} else {
+			logrus.WithFields(logrus.Fields{
+				"id": id,
+				"rank": rank,
+			}).Error(err.Error())
+		}
+		user, exist := Users[id]
+		if exist {
+			if user.Rank != rank {
+				logrus.Info("New rank (", rank, ") for user ", id)
+				if user.CurrentRole != "" {
+					logrus.Info("Remove role for ", id)
+					removeRole(discordClient, guild, id, user.CurrentRole, dryrun)
+					Users[id] = &UserRole{Rank: rank, CurrentRole: ""}
+				}
+				if rank == 1 {
+					logrus.Info("Add role `top1` for ", id)
+					addRole(discordClient, guild, id, rolesID["top1"], dryrun)
+					Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top1"]}
+				} else if rank <= 10 {
+					logrus.Info("Add role `top10` for ", id)
+					addRole(discordClient, guild, id, rolesID["top10"], dryrun)
+					Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top10"]}
+				} else if rank <= 50 {
+					logrus.Info("Add role `top50` for ", id)
+					addRole(discordClient, guild, id, rolesID["top50"], dryrun)
+					Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top50"]}
+				}
+			}
+		}
+	}
+	return nil
+}
+
+func forceUpdate(discordClient *discordgo.Session, dryrun bool) (error) {
+	APIResponse, err := getUsers42CTF()
+	if err != nil {
+		return err
+	}
+	for id, rank := range APIResponse {
+		tmp, err := discordClient.GuildMember(guild, id)
+		if err == nil {
+			if slices.Contains(tmp.Roles, rolesID["top1"]) {
+				removeRole(discordClient, guild, id, rolesID["top1"], dryrun)
+			}
+			if slices.Contains(tmp.Roles, rolesID["top10"]) {
+				removeRole(discordClient, guild, id, rolesID["top10"], dryrun)
+			}
+			if slices.Contains(tmp.Roles, rolesID["top50"]) {
+				removeRole(discordClient, guild, id, rolesID["top50"], dryrun)
+			}
+			Users[id] = &UserRole{Rank: rank, CurrentRole: ""}
+		} else {
+			logrus.Error("ID not found", id)
+		}
+		_, exist := Users[id]
+		if exist {
+			if rank == 1 {
+				logrus.Info("Add role `top1` for ", id)
+				addRole(discordClient, guild, id, rolesID["top1"], dryrun)
+				Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top1"]}
+			} else if rank <= 10 {
+				logrus.Info("Add role `top10` for ", id)
+				addRole(discordClient, guild, id, rolesID["top10"], dryrun)
+				Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top10"]}
+			} else if rank <= 50 {
+				logrus.Info("Add role `top50` for ", id)
+				addRole(discordClient, guild, id, rolesID["top50"], dryrun)
+				Users[id] = &UserRole{Rank: rank, CurrentRole: rolesID["top50"]}
+			}
+		}
+	}
+	return nil
+}

src/Dockerfile

@@ -1,17 +0,0 @@
-FROM debian:9
-
-ADD . /app/
-WORKDIR /app
-
-RUN	apt-get update -yq											\
-	&& apt-get install libffi-dev								\
-	&& apt-get install libnacl-dev								\
-	&& apt-get install python3-dev								\
-	&& apt-get install python3-pip								\
-	&& python3 -m pip install -U discord.py						\
-	&& python3 -m pip install -U asyncio						\
-	&& python3 -m pip install -U requests						\
-
-
-CMD python3 bot.py
-

src/bot.py

@@ -1,77 +0,0 @@
-import discord
-import asyncio
-import time
-import requests
-import colorama
-
-intents = discord.Intents.all()
-
-client = discord.Client(intents=intents)
-
-def logger(type, message):
-	color = colorama.Fore.WHITE
-	if type == 'info':
-		color = colorama.Fore.GREEN
-	if type == 'warn':
-		color = colorama.Fore.YELLOW
-	if type == 'error':
-		color = colorama.Fore.RED
-	print(f"{colorama.Fore.BLUE}[{color}{type.upper()}{colorama.Fore.BLUE}]{colorama.Fore.WHITE}{message}")
-
-async def get_json():
-	url = 'https://preprod.42ctf.org/api/bot/discord?token=test'
-	resp = requests.get(url=url)
-	data = resp.json()
-	return data
-
-async def put_role(guild, userid, rank):
-	logger("info", f"userid {userid} has rank {rank}")
-	member = guild.get_member(user_id=userid)
-	if not member:
-		logger("warn", f"member with userid {userid} not found")
-		return
-	if rank < 1: # if rank is not in the top 50 or invalid
-		logger("error", f"rank {rank} is invalid")
-		return
-	if rank > 50:
-		await member.remove_roles(guild.get_role(role_id=943545946309029958))
-		await member.remove_roles(guild.get_role(role_id=943545892030517278))
-		await member.remove_roles(guild.get_role(role_id=943545814758854686))
-		return
-	if rank > 10: # if rank is not in the top 10
-		role = guild.get_role(role_id=943545946309029958)
-		await member.remove_roles(guild.get_role(role_id=943545892030517278))
-		await member.remove_roles(guild.get_role(role_id=943545814758854686))
-		await member.add_roles(role)
-		return
-	if rank > 1: # if rank is not top 1
-		role = guild.get_role(role_id=943545892030517278)
-		await member.remove_roles(guild.get_role(role_id=943545814758854686))
-		await member.remove_roles(guild.get_role(role_id=943545946309029958))
-		await member.add_roles(role)
-		return
-	if rank == 1: # if rank is top 1
-		role = guild.get_role(role_id=943545814758854686)
-		await member.remove_roles(guild.get_role(role_id=943545946309029958))
-		await member.remove_roles(guild.get_role(role_id=943545892030517278))
-		await member.add_roles(role)
-		return
-	
-
-async def routine(guild):
-	logger("info", "routine started")
-	while 1:
-		data = await get_json()
-		for key in data:
-			await put_role(guild, int(key), int(data[key]))
-		time.sleep(5000)
-
-@client.event
-async def on_ready():
-	logger("info", "logged as " + client.user.name)
-	await client.change_presence(activity=discord.Activity(type=discord.ActivityType.playing, name='etre refait en python'))
-	guild = client.get_guild(943459216901955604)
-	await routine(guild)
-
-client.run('OTQzNTA2NTgxMjkzNzc2OTQ3.GQ2l77.IyrlZv8-ESoXfL1MfSQw_ec3sXqOV1eh9PdNh4')
-