update: add setliveimage command

Allow for custom images in live embeds

internal/bot/bot.go

@@ -84,7 +84,7 @@ func (b *Bot) monitorUsers() {
 			}
 			defer tx.Rollback()
 
-			rows, err := tx.Query("SELECT guild_id, user_id, username, notification_channel, last_post_id, last_stream_start, mention_role, avatar_location, avatar_location_updated_at FROM monitored_users")
+			rows, err := tx.Query("SELECT guild_id, user_id, username, notification_channel, last_post_id, last_stream_start, mention_role, avatar_location, avatar_location_updated_at, live_image_url FROM monitored_users")
 			if err != nil {
 				return err
 			}
@@ -101,9 +101,10 @@ func (b *Bot) monitorUsers() {
 					MentionRole             string
 					AvatarLocation          string
 					AvatarLocationUpdatedAt int64
+					LiveImageURL            string
 				}
 
-				err := rows.Scan(&user.GuildID, &user.UserID, &user.Username, &user.NotificationChannel, &user.LastPostID, &user.LastStreamStart, &user.MentionRole, &user.AvatarLocation, &user.AvatarLocationUpdatedAt)
+				err := rows.Scan(&user.GuildID, &user.UserID, &user.Username, &user.NotificationChannel, &user.LastPostID, &user.LastStreamStart, &user.MentionRole, &user.AvatarLocation, &user.AvatarLocationUpdatedAt, &user.LiveImageURL)
 				if err != nil {
 					log.Printf("Error scanning row: %v", err)
 					continue
@@ -145,7 +146,7 @@ func (b *Bot) monitorUsers() {
 						return err
 					}
 
-					embedMsg := embed.CreateLiveStreamEmbed(user.Username, streamInfo, user.AvatarLocation)
+					embedMsg := embed.CreateLiveStreamEmbed(user.Username, streamInfo, user.AvatarLocation, user.LiveImageURL)
 
 					mention := "@everyone"
 					if user.MentionRole != "" {

internal/bot/commands.go

@@ -47,6 +47,24 @@ func (b *Bot) registerCommands() {
 			Name:        "list",
 			Description: "List all monitored models",
 		},
+		{
+			Name:        "setliveimage",
+			Description: "Set a custom live image for a model",
+			Options: []*discordgo.ApplicationCommandOption{
+				{
+					Type:        discordgo.ApplicationCommandOptionString,
+					Name:        "username",
+					Description: "The username of the model",
+					Required:    true,
+				},
+				{
+					Type:        discordgo.ApplicationCommandOptionAttachment,
+					Name:        "image",
+					Description: "The image to use for live notifications",
+					Required:    true,
+				},
+			},
+		},
 	}
 
 	_, err := b.Session.ApplicationCommandBulkOverwrite(b.Session.State.User.ID, "", commands)

internal/bot/handlers.go

@@ -28,6 +28,8 @@ func (b *Bot) interactionCreate(s *discordgo.Session, i *discordgo.InteractionCr
 		b.handleRemoveCommand(s, i)
 	case "list":
 		b.handleListCommand(s, i)
+	case "setliveimage":
+		b.handleSetLiveImageCommand(s, i)
 	}
 }
 
@@ -94,9 +96,9 @@ func (b *Bot) handleAddCommand(s *discordgo.Session, i *discordgo.InteractionCre
 	// Store the monitored user in the database
 	_, err = b.DB.Exec(`
         INSERT OR REPLACE INTO monitored_users 
-        (guild_id, user_id, username, notification_channel, last_post_id, last_stream_start, mention_role, avatar_location, avatar_location_updated_at) 
-        VALUES (?, ?, ?, ?, '', 0, ?, ?, ?)
-    `, i.GuildID, accountInfo.ID, username, channel.ID, mentionRole, avatarLocation, time.Now().Unix())
+        (guild_id, user_id, username, notification_channel, last_post_id, last_stream_start, mention_role, avatar_location, avatar_location_updated_at, live_image_url) 
+        VALUES (?, ?, ?, ?, '', 0, ?, ?, ?, ?)
+		`, i.GuildID, accountInfo.ID, username, channel.ID, mentionRole, avatarLocation, time.Now().Unix(), "")
 	if err != nil {
 		b.respondToInteraction(s, i, fmt.Sprintf("Error storing user: %v", err))
 		return
@@ -182,6 +184,64 @@ func (b *Bot) handleListCommand(s *discordgo.Session, i *discordgo.InteractionCr
 	b.respondToInteraction(s, i, response)
 }
 
+func (b *Bot) handleSetLiveImageCommand(s *discordgo.Session, i *discordgo.InteractionCreate) {
+	// Acknowledge the interaction immediately
+	err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
+		Type: discordgo.InteractionResponseDeferredChannelMessageWithSource,
+	})
+	if err != nil {
+		log.Printf("Error acknowledging interaction: %v", err)
+		return
+	}
+
+	if !b.hasAdminOrModPermissions(s, i) {
+		b.editInteractionResponse(s, i, "You need administrator permissions to use this command.")
+		return
+	}
+
+	options := i.ApplicationCommandData().Options
+	username := options[0].StringValue()
+
+	var imageURL string
+	if len(i.ApplicationCommandData().Resolved.Attachments) > 0 {
+		for _, attachment := range i.ApplicationCommandData().Resolved.Attachments {
+			imageURL = attachment.URL
+			//log.Println(imageURL)
+			break
+		}
+	}
+
+	if imageURL == "" {
+		b.editInteractionResponse(s, i, "Please attach an image to set as the live image.")
+		return
+	}
+
+	// Update the database with the new live image URL
+	_, err = b.DB.Exec(`
+        UPDATE monitored_users 
+        SET live_image_url = ? 
+        WHERE guild_id = ? AND username = ?
+    `, imageURL, i.GuildID, username)
+
+	if err != nil {
+		log.Printf("Error updating live image URL: %v", err)
+		b.editInteractionResponse(s, i, "An error occurred while setting the live image.")
+		return
+	}
+
+	b.editInteractionResponse(s, i, fmt.Sprintf("Live image for %s has been set successfully.", username))
+}
+
+// Add this new helper function
+func (b *Bot) editInteractionResponse(s *discordgo.Session, i *discordgo.InteractionCreate, content string) {
+	_, err := s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{
+		Content: &content,
+	})
+	if err != nil {
+		log.Printf("Error editing interaction response: %v", err)
+	}
+}
+
 func (b *Bot) hasAdminOrModPermissions(s *discordgo.Session, i *discordgo.InteractionCreate) bool {
 	// Check if the user is the server owner
 	guild, err := s.Guild(i.GuildID)

internal/database/databse.go

@@ -35,6 +35,7 @@ func createTables() {
 			mention_role TEXT,
 			avatar_location TEXT,
 			avatar_location_updated_at INTEGER,
+			live_image_url	TEXT,
 			PRIMARY KEY (guild_id, user_id)
 		)
 	`)

internal/embed/embed.go

@@ -9,7 +9,7 @@ import (
 	"github.com/fvckgrimm/discord-fansly-notify/api"
 )
 
-func CreateLiveStreamEmbed(username string, streamInfo *api.StreamResponse, avatarLocation string) *discordgo.MessageEmbed {
+func CreateLiveStreamEmbed(username string, streamInfo *api.StreamResponse, avatarLocation string, liveImageURL string) *discordgo.MessageEmbed {
 	liveURL := fmt.Sprintf("https://fansly.com/live/%s", username)
 	creatorUrl := fmt.Sprintf("https://fansly.com/%s", username)
 
@@ -40,6 +40,13 @@ func CreateLiveStreamEmbed(username string, streamInfo *api.StreamResponse, avat
 		},
 		Timestamp: time.Now().Format(time.RFC3339),
 	}
+
+	if liveImageURL != "" {
+		embed.Image = &discordgo.MessageEmbedImage{
+			URL: liveImageURL,
+		}
+	}
+
 	return embed
 }
 

migrate.py

@@ -1,10 +1,9 @@
 import sqlite3
-import time
+# import time
 
 # Connect to the old and new databases
 old_conn = sqlite3.connect("bot-old.db")
 new_conn = sqlite3.connect("bot.db")
-
 old_cursor = old_conn.cursor()
 new_cursor = new_conn.cursor()
 
@@ -20,6 +19,7 @@
     mention_role TEXT,
     avatar_location TEXT,
     avatar_location_updated_at INTEGER,
+    live_image_url TEXT,
     PRIMARY KEY (guild_id, user_id)
 )
 """)
@@ -30,16 +30,17 @@
 
 # Insert the old records into the new table
 for record in old_records:
-    # Assuming the order of columns in the old table is:
-    # guild_id, user_id, username, notification_channel, last_post_id, last_stream_start, mention_role, avatar_location
     new_cursor.execute(
         """
     INSERT INTO monitored_users 
-    (guild_id, user_id, username, notification_channel, last_post_id, last_stream_start, mention_role, avatar_location, avatar_location_updated_at)
-    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+    (guild_id, user_id, username, notification_channel, last_post_id, last_stream_start, mention_role, avatar_location, avatar_location_updated_at, live_image_url)
+    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
     """,
-        (*record, int(time.time())),
-    )  # Add current timestamp for avatar_location_updated_at
+        (
+            *record,
+            None,
+        ),  # Add all existing fields from old record, plus None for live_image_url
+    )
 
 # Commit the changes and close the connections
 new_conn.commit()