William P
a6dc4e27b9
All checks were successful
Build only (for PRs) / build-only (pull_request) Successful in 3m36s
559 lines
16 KiB
Go
559 lines
16 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"sort"
|
|
"strings"
|
|
"syscall"
|
|
|
|
"github.com/bwmarrin/discordgo"
|
|
)
|
|
|
|
func main() {
|
|
out_dir := os.Getenv("OUT_PATH")
|
|
temp_dir := os.Getenv("TEMP_PATH")
|
|
bot_token := os.Getenv("DISCORD_TOKEN")
|
|
|
|
if out_dir == "" {
|
|
panic("No output dir specified")
|
|
}
|
|
|
|
if temp_dir == "" {
|
|
panic("No temp dir specified")
|
|
}
|
|
|
|
s, err := discordgo.New("Bot " + bot_token)
|
|
if err != nil {
|
|
log.Fatalf("Invalid bot parameters: %v", err)
|
|
}
|
|
|
|
err = s.Open()
|
|
if err != nil {
|
|
log.Fatalf("Error opening connection: %v", err)
|
|
}
|
|
|
|
var defaultMemberPermissions int64 = discordgo.PermissionAllText
|
|
var interactionPrivateChannel = discordgo.InteractionContextPrivateChannel
|
|
|
|
var commands = []*discordgo.ApplicationCommand{
|
|
{
|
|
Name: "download",
|
|
Description: "Download video and save it to 'youtube-vids",
|
|
DefaultMemberPermissions: &defaultMemberPermissions,
|
|
Contexts: &[]discordgo.InteractionContextType{interactionPrivateChannel},
|
|
Options: []*discordgo.ApplicationCommandOption{
|
|
{
|
|
Name: "url",
|
|
Description: "URL",
|
|
Type: discordgo.ApplicationCommandOptionString,
|
|
Required: true,
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Name: "version",
|
|
Description: "Show application version",
|
|
DefaultMemberPermissions: &defaultMemberPermissions,
|
|
Contexts: &[]discordgo.InteractionContextType{interactionPrivateChannel},
|
|
},
|
|
{
|
|
Name: "download video",
|
|
Type: discordgo.MessageApplicationCommand,
|
|
IntegrationTypes: &[]discordgo.ApplicationIntegrationType{
|
|
discordgo.ApplicationIntegrationUserInstall,
|
|
},
|
|
|
|
Contexts: &[]discordgo.InteractionContextType{
|
|
discordgo.InteractionContextBotDM,
|
|
discordgo.InteractionContextPrivateChannel,
|
|
},
|
|
},
|
|
}
|
|
|
|
var componentHandlers = map[string]func(s *discordgo.Session, i *discordgo.InteractionCreate){
|
|
"video_select": func(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
|
// Get interaction state with fallback search
|
|
var state *InteractionState
|
|
|
|
// First try current token
|
|
state = getInteractionState(i.Interaction.Token)
|
|
|
|
// If not found, try to find state associated with this interaction context
|
|
if state == nil {
|
|
interactionStatesMutex.RLock()
|
|
for token, s := range interactionStates {
|
|
if s != nil && s.FormatOptions != nil {
|
|
// Found a state, transfer it to current token
|
|
state = s
|
|
// Copy to current token and clean up old one
|
|
interactionStatesMutex.RUnlock()
|
|
setInteractionState(i.Interaction.Token, state)
|
|
deleteInteractionState(token)
|
|
break
|
|
}
|
|
}
|
|
if state == nil {
|
|
interactionStatesMutex.RUnlock()
|
|
}
|
|
}
|
|
if state == nil || state.FormatOptions == nil {
|
|
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "Error: Session expired. Please start over with /download.",
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
// Get requester user ID
|
|
state.Requester = i.User.ID
|
|
|
|
// Get selected video format
|
|
selectedValues := i.MessageComponentData().Values
|
|
if len(selectedValues) == 0 {
|
|
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "Error: No video format selected",
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
// Store selected video format ID and name
|
|
state.VideoFormatID = selectedValues[0]
|
|
for _, vOpt := range state.FormatOptions.VideoOptions {
|
|
if vOpt.FormatID == selectedValues[0] {
|
|
label := fmt.Sprintf("%s (%s", vOpt.Resolution, vOpt.Ext)
|
|
if vOpt.TBR != nil {
|
|
label += fmt.Sprintf(", %.0fkbps", *vOpt.TBR)
|
|
}
|
|
label += ")"
|
|
state.VideoFormatName = label
|
|
break
|
|
}
|
|
}
|
|
setInteractionState(i.Interaction.Token, state)
|
|
|
|
// Build audio format options
|
|
audioMenuOptions := buildAudioMenuOptions(state.FormatOptions.AudioOptions)
|
|
|
|
// Build video format options (to keep them visible but disabled)
|
|
videoMenuOptions := buildVideoMenuOptions(state.FormatOptions.VideoOptions)
|
|
|
|
// Update components - disable video select, enable audio select
|
|
updatedComponents := []discordgo.MessageComponent{
|
|
discordgo.ActionsRow{
|
|
Components: []discordgo.MessageComponent{
|
|
discordgo.SelectMenu{
|
|
CustomID: "video_select",
|
|
Placeholder: "Choose a video format...",
|
|
MaxValues: 1,
|
|
Disabled: true,
|
|
Options: videoMenuOptions,
|
|
},
|
|
},
|
|
},
|
|
discordgo.ActionsRow{
|
|
Components: []discordgo.MessageComponent{
|
|
discordgo.SelectMenu{
|
|
CustomID: "audio_select",
|
|
Placeholder: "Choose an audio format...",
|
|
MaxValues: 1,
|
|
Disabled: false,
|
|
Options: audioMenuOptions,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseUpdateMessage,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "Now select an audio format:",
|
|
Components: updatedComponents,
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error: %v", err)
|
|
}
|
|
},
|
|
"audio_select": func(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
|
// Get URL from interaction state with fallback search
|
|
var state *InteractionState
|
|
|
|
// First try current token
|
|
state = getInteractionState(i.Interaction.Token)
|
|
|
|
// If not found, try to find state associated with this interaction context
|
|
if state == nil {
|
|
interactionStatesMutex.RLock()
|
|
for token, s := range interactionStates {
|
|
if s != nil && s.URL != "" {
|
|
// Found a state, transfer it to current token
|
|
state = s
|
|
// Copy to current token and clean up old one
|
|
interactionStatesMutex.RUnlock()
|
|
setInteractionState(i.Interaction.Token, state)
|
|
deleteInteractionState(token)
|
|
break
|
|
}
|
|
}
|
|
if state == nil {
|
|
interactionStatesMutex.RUnlock()
|
|
}
|
|
}
|
|
if state == nil {
|
|
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseUpdateMessage,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "Error: No video URL found. Please start over with /download.",
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
// Get selected audio format
|
|
selectedValues := i.MessageComponentData().Values
|
|
audioFormatID := ""
|
|
if len(selectedValues) > 0 {
|
|
audioFormatID = selectedValues[0]
|
|
}
|
|
|
|
// Store selected audio format ID and name
|
|
state.AudioFormatID = audioFormatID
|
|
for _, aOpt := range state.FormatOptions.AudioOptions {
|
|
if aOpt.FormatID == audioFormatID {
|
|
label := aOpt.Format
|
|
if aOpt.Language != nil {
|
|
label += fmt.Sprintf(" [%s]", *aOpt.Language)
|
|
}
|
|
if aOpt.TBR != nil {
|
|
label += fmt.Sprintf(" (%.0fkbps)", *aOpt.TBR)
|
|
}
|
|
state.AudioFormatName = label
|
|
break
|
|
}
|
|
}
|
|
|
|
response := ""
|
|
if state.URL != "" {
|
|
// Respond immediately to prevent timeout
|
|
response = fmt.Sprintf("%s **Starting download**", loading_emoji)
|
|
|
|
// Clean up state before responding
|
|
deleteInteractionState(i.Interaction.Token)
|
|
} else {
|
|
response = "I don't see a video here :("
|
|
}
|
|
err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseUpdateMessage,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: response,
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error: %v", err)
|
|
}
|
|
if state.URL != "" {
|
|
go startAsyncDownload(s, i, state.Requester, state.URL, state.VideoFormatID, state.VideoFormatName, state.AudioFormatID, state.AudioFormatName, out_dir, temp_dir)
|
|
}
|
|
},
|
|
}
|
|
|
|
var commandHandlers = map[string]func(s *discordgo.Session, i *discordgo.InteractionCreate){
|
|
"download": func(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
|
options := i.ApplicationCommandData().Options
|
|
optionMap := make(map[string]*discordgo.ApplicationCommandInteractionDataOption, len(options))
|
|
for _, opt := range options {
|
|
optionMap[opt.Name] = opt
|
|
}
|
|
|
|
var url string
|
|
if option, ok := optionMap["url"]; ok {
|
|
url = option.StringValue()
|
|
} else {
|
|
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "Error: No URL provided",
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
// Send initial "fetching formats" response
|
|
err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: fmt.Sprintf("%s Fetching available formats...", loading_emoji),
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error: %v", err)
|
|
return
|
|
}
|
|
|
|
// Fetch formats asynchronously
|
|
go fetchAndShowFormats(s, i, url)
|
|
},
|
|
"download video": func(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
|
data := i.ApplicationCommandData()
|
|
targetMsg, ok := data.Resolved.Messages[data.TargetID]
|
|
if !ok || targetMsg == nil {
|
|
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "Error: Could not find the target message",
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
return
|
|
}
|
|
url := extractURLFromString(targetMsg.Content)
|
|
if url == "" {
|
|
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "Error: No URL provided",
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
// Send initial "fetching formats" response
|
|
err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: fmt.Sprintf("%s Fetching available formats...", loading_emoji),
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error: %v", err)
|
|
return
|
|
}
|
|
|
|
// Fetch formats asynchronously
|
|
go fetchAndShowFormats(s, i, url)
|
|
},
|
|
"version": func(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
|
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "[yt-dlp-bot](https://git.dubyatp.xyz/williamp/yt-dlp-bot) by dubyatp",
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
},
|
|
}
|
|
|
|
s.AddHandler(func(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
|
switch i.Type {
|
|
case discordgo.InteractionApplicationCommand:
|
|
if h, ok := commandHandlers[i.ApplicationCommandData().Name]; ok {
|
|
h(s, i)
|
|
}
|
|
case discordgo.InteractionMessageComponent:
|
|
customID := i.MessageComponentData().CustomID
|
|
if h, ok := componentHandlers[customID]; ok {
|
|
h(s, i)
|
|
} else if strings.HasPrefix(customID, "retry:") {
|
|
url := strings.TrimPrefix(customID, "retry:")
|
|
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseUpdateMessage,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: fmt.Sprintf("%s Fetching available formats...", loading_emoji),
|
|
Components: []discordgo.MessageComponent{},
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
go fetchAndShowFormats(s, i, url)
|
|
}
|
|
}
|
|
})
|
|
|
|
log.Println("Initialize loading emoji")
|
|
initLoadingEmoji(s)
|
|
|
|
|
|
log.Println("Adding commands")
|
|
registeredCommands := make([]*discordgo.ApplicationCommand, len(commands))
|
|
for i, v := range commands {
|
|
cmd, err := s.ApplicationCommandCreate(s.State.User.ID, "", v)
|
|
if err != nil {
|
|
log.Panicf("Cannot create '%v' command: %v", v.Name, err)
|
|
}
|
|
registeredCommands[i] = cmd
|
|
}
|
|
|
|
log.Println("Bot is now running. Press CTRL+C to exit")
|
|
sc := make(chan os.Signal, 1)
|
|
signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
|
|
<-sc
|
|
|
|
s.Close()
|
|
}
|
|
|
|
func buildVideoMenuOptions(videoOptions []VideoOption) []discordgo.SelectMenuOption {
|
|
opts := make([]discordgo.SelectMenuOption, 0, 25)
|
|
for _, vOpt := range videoOptions {
|
|
label := fmt.Sprintf("%s (%s", vOpt.Resolution, vOpt.Ext)
|
|
if vOpt.TBR != nil {
|
|
label += fmt.Sprintf(", %.0fkbps", *vOpt.TBR)
|
|
}
|
|
label += ")"
|
|
if len(label) > 100 {
|
|
label = label[:97] + "..."
|
|
}
|
|
opts = append(opts, discordgo.SelectMenuOption{
|
|
Label: label,
|
|
Value: vOpt.FormatID,
|
|
})
|
|
if len(opts) >= 25 {
|
|
break
|
|
}
|
|
}
|
|
return opts
|
|
}
|
|
|
|
func buildAudioMenuOptions(audioOptions []AudioOption) []discordgo.SelectMenuOption {
|
|
opts := make([]discordgo.SelectMenuOption, 0, 25)
|
|
for _, aOpt := range audioOptions {
|
|
label := aOpt.Format
|
|
if aOpt.Language != nil {
|
|
label += fmt.Sprintf(" [%s]", *aOpt.Language)
|
|
}
|
|
if aOpt.TBR != nil {
|
|
label += fmt.Sprintf(" (%.0fkbps)", *aOpt.TBR)
|
|
}
|
|
if len(label) > 100 {
|
|
label = label[:97] + "..."
|
|
}
|
|
opts = append(opts, discordgo.SelectMenuOption{
|
|
Label: label,
|
|
Value: aOpt.FormatID,
|
|
})
|
|
if len(opts) >= 25 {
|
|
break
|
|
}
|
|
}
|
|
return opts
|
|
}
|
|
|
|
func fetchAndShowFormats(s *discordgo.Session, i *discordgo.InteractionCreate, url string) {
|
|
formatOptions, err := GetFormats(url)
|
|
if err != nil {
|
|
_, err = s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{
|
|
Content: ptr("❌ Error fetching formats: " + err.Error()),
|
|
Components: &[]discordgo.MessageComponent{
|
|
&discordgo.ActionsRow{
|
|
Components: []discordgo.MessageComponent{
|
|
discordgo.Button{
|
|
Label: "Retry",
|
|
CustomID: fmt.Sprintf("retry:%s", url),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error updating interaction: %v", err)
|
|
}
|
|
return
|
|
}
|
|
|
|
// Sort video formats: highest resolution first, then by bitrate
|
|
sort.Slice(formatOptions.VideoOptions, func(x, y int) bool {
|
|
heightX, heightY := 0, 0
|
|
if formatOptions.VideoOptions[x].Height != nil {
|
|
heightX = *formatOptions.VideoOptions[x].Height
|
|
}
|
|
if formatOptions.VideoOptions[y].Height != nil {
|
|
heightY = *formatOptions.VideoOptions[y].Height
|
|
}
|
|
if heightX != heightY {
|
|
return heightX > heightY
|
|
}
|
|
tbrX, tbrY := 0.0, 0.0
|
|
if formatOptions.VideoOptions[x].TBR != nil {
|
|
tbrX = *formatOptions.VideoOptions[x].TBR
|
|
}
|
|
if formatOptions.VideoOptions[y].TBR != nil {
|
|
tbrY = *formatOptions.VideoOptions[y].TBR
|
|
}
|
|
return tbrX > tbrY
|
|
})
|
|
|
|
// Sort audio formats: highest bitrate first
|
|
sort.Slice(formatOptions.AudioOptions, func(x, y int) bool {
|
|
tbrX, tbrY := 0.0, 0.0
|
|
if formatOptions.AudioOptions[x].TBR != nil {
|
|
tbrX = *formatOptions.AudioOptions[x].TBR
|
|
}
|
|
if formatOptions.AudioOptions[y].TBR != nil {
|
|
tbrY = *formatOptions.AudioOptions[y].TBR
|
|
}
|
|
return tbrX > tbrY
|
|
})
|
|
|
|
videoMenuOptions := buildVideoMenuOptions(formatOptions.VideoOptions)
|
|
audioMenuOptions := buildAudioMenuOptions(formatOptions.AudioOptions)
|
|
|
|
if len(videoMenuOptions) == 0 || len(audioMenuOptions) == 0 {
|
|
_, err = s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{
|
|
Content: ptr("❌ No separate video/audio streams found for this URL. The source may only provide combined formats."),
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error updating interaction: %v", err)
|
|
}
|
|
return
|
|
}
|
|
|
|
setInteractionState(i.Interaction.Token, &InteractionState{
|
|
URL: url,
|
|
FormatOptions: formatOptions,
|
|
})
|
|
|
|
_, err = s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{
|
|
Content: ptr("Select a video format:"),
|
|
Components: &[]discordgo.MessageComponent{
|
|
discordgo.ActionsRow{
|
|
Components: []discordgo.MessageComponent{
|
|
discordgo.SelectMenu{
|
|
CustomID: "video_select",
|
|
Placeholder: "Choose a video format...",
|
|
MaxValues: 1,
|
|
Options: videoMenuOptions,
|
|
},
|
|
},
|
|
},
|
|
discordgo.ActionsRow{
|
|
Components: []discordgo.MessageComponent{
|
|
discordgo.SelectMenu{
|
|
CustomID: "audio_select",
|
|
Placeholder: "Choose an audio format...",
|
|
MaxValues: 1,
|
|
Disabled: true,
|
|
Options: audioMenuOptions,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error updating interaction: %v", err)
|
|
}
|
|
}
|