diff --git a/database.go b/database.go index bbbc7e1..e93f93b 100644 --- a/database.go +++ b/database.go @@ -8,6 +8,7 @@ import ( const videoTable = "CREATE TABLE IF NOT EXISTS `video` ( `filename` VARCHAR(15) NOT NULL, `video` TEXT NOT NULL, PRIMARY KEY (`filename`) )" const albumTable = "CREATE TABLE IF NOT EXISTS `album` ( `filename` VARCHAR(11) NOT NULL, `album` TEXT NOT NULL, PRIMARY KEY (`filename`) )" +const audioTable = "CREATE TABLE IF NOT EXISTS `audio` ( `filename` VARCHAR(15) NOT NULL, `audio` TEXT NOT NULL, PRIMARY KEY (`filename`) )" func execQuery(query ...string) { for _, q := range query { @@ -25,6 +26,7 @@ func load() { bytes []byte ) + // Videos rows, err := db.Query("SELECT filename, video FROM video") if err != nil { lit.Error("Error executing query, %s", err) @@ -44,6 +46,7 @@ func load() { cacheVideo[filename] = &video } + // Photos rows, err = db.Query("SELECT filename, album FROM album") if err != nil { lit.Error("Error executing query, %s", err) @@ -62,6 +65,26 @@ func load() { err = json.Unmarshal(bytes, &photos) cacheAlbum[filename] = &photos } + + // Audio + rows, err = db.Query("SELECT filename, audio FROM audio") + if err != nil { + lit.Error("Error executing query, %s", err) + return + } + + for rows.Next() { + var audio tele.Audio + + err = rows.Scan(&filename, &bytes) + if err != nil { + lit.Error("Error scanning row, %s", err) + continue + } + + err = json.Unmarshal(bytes, &audio) + cacheAudio[filename] = &audio + } } func saveVideo(video *tele.Video) { @@ -83,3 +106,13 @@ func saveAlbum(album *[]*tele.Photo, filename string) { return } } + +func saveAudio(audio *tele.Audio) { + bytes, _ := json.Marshal(audio) + + _, err := db.Exec("INSERT INTO audio (filename, audio) VALUES (?, ?)", audio.FileName, bytes) + if err != nil { + lit.Error("Error executing query, %s", err) + return + } +} diff --git a/events.go b/events.go index f188801..df7498f 100644 --- a/events.go +++ b/events.go @@ -12,56 +12,57 @@ func videoDownload(c tele.Context) error { link := cleanURL(c.Message().EntityText(e)) if contains(link, cfg.URLs) { - // Use the downloader to get videos and albums from tiktok + var ( + media Media + filename string + hit bool + ) + if strings.Contains(link, "tiktok.com") { - filename, hit, media := downloadTikTok(link) - if filename != "" { - if media == Video { - err := c.Reply(cacheVideo[filename], tele.Silent) - if err == nil { - if !hit { - go saveVideo(cacheVideo[filename]) - continue - } - } - } else { - if _, ok := cacheAlbum[filename]; ok { - var err error - - photos := *cacheAlbum[filename] - album := make(tele.Album, 0, 10) - - for i := 0; i < len(photos); i += 10 { - // Add photos to album - for j := 0; j < 10; j++ { - if i+j < len(photos) { - album = append(album, photos[i+j]) - } - } + // Use the downloader to get videos and albums from TikTok + filename, hit, media = downloadTikTok(link) + } else { + filename, hit = downloadYtDlp(link) + media = Video + } - err = c.SendAlbum(album, tele.Silent) - if err != nil { - break + if filename != "" { + if media == Video { + err := c.Reply(cacheVideo[filename], tele.Silent) + if err == nil && !hit { + go saveVideo(cacheVideo[filename]) + } + } else { + if _, ok := cacheAlbum[filename]; ok { + photos := *cacheAlbum[filename] + album := make(tele.Album, 0, 10) + + for i := 0; i < len(photos); i += 10 { + // Add photos to album + for j := 0; j < 10; j++ { + if i+j < len(photos) { + album = append(album, photos[i+j]) } } - if err == nil { - if !hit { - go saveAlbum(&photos, filename) - continue - } + _ = c.SendAlbum(album, tele.Silent) + } + + if !hit { + go saveAlbum(&photos, filename) + } + + // Handle audio + filename, hit = downloadAudio(link) + if filename != "" { + err := c.Reply(cacheAudio[filename], tele.Silent) + if err == nil && !hit { + go saveAudio(cacheAudio[filename]) } } } } } - - filename, hit := downloadYtDlp(link) - - err := c.Reply(cacheVideo[filename], tele.Silent) - if err == nil && !hit { - go saveVideo(cacheVideo[filename]) - } } else { // For twitter, we send the same url with only fx appended to it if strings.HasPrefix(link, "https://twitter.com") { @@ -86,42 +87,74 @@ func videoDownload(c tele.Context) error { func inlineQuery(c tele.Context) error { var ( - results = make([]tele.Result, 1) text = c.Query().Text + results = make([]tele.Result, 0, 1) ) if isValidURL(text) && contains(text, cfg.URLs) { + var ( + media Media + filename string + hit bool + ) + text = cleanURL(text) - filename, hit := downloadYtDlp(text) - // Upload video to channel, so we can send it even in inline mode - _, err := c.Bot().Send(tele.ChatID(cfg.Channel), cacheVideo[filename]) - if err != nil { - return err + if strings.Contains(text, "tiktok.com") { + // Use the downloader to get videos and albums from TikTok + filename, hit, media = downloadTikTok(text) + } else { + filename, hit = downloadYtDlp(text) + media = Video } - if !hit { - go saveVideo(cacheVideo[filename]) - } + if filename != "" { + if media == Video { + // Upload video to channel, so we can send it even in inline mode + _, err := c.Bot().Send(tele.ChatID(cfg.Channel), cacheVideo[filename]) + if err != nil { + return err + } - // Create result - results[0] = &tele.VideoResult{ - Cache: cacheVideo[filename].FileID, - Title: "Send video", - MIME: "video/mp4", - } + if !hit { + go saveVideo(cacheVideo[filename]) + } + + // Create result + results = append(results, &tele.VideoResult{ + Cache: cacheVideo[filename].FileID, + Title: "Send video", + MIME: "video/mp4", + }) + + results[0].SetResultID(filename) + } else { + if _, ok := cacheAlbum[filename]; ok { + photos := *cacheAlbum[filename] + + for i, p := range photos { + results = append(results, &tele.PhotoResult{ + URL: p.FileURL, + ThumbURL: p.FileURL, + }) + results[i].SetResultID(idGen(p.FileURL)) + } - results[0].SetResultID(filename) + if !hit { + go saveAlbum(&photos, filename) + } + } + } + } - // Send video return c.Answer(&tele.QueryResponse{ Results: results, CacheTime: 86400, // one day }) } else { - results[0] = &tele.ArticleResult{ - Title: "Not a valid instagram URL!", - } + results = append(results, &tele.ArticleResult{ + Title: "Not a valid URL!", + }) results[0].SetResultID(text) diff --git a/main.go b/main.go index 9a7f4c1..c35b8ef 100644 --- a/main.go +++ b/main.go @@ -16,6 +16,7 @@ var ( cfg config cacheVideo map[string]*tele.Video cacheAlbum map[string]*[]*tele.Photo + cacheAudio map[string]*tele.Audio db *sql.DB ) @@ -42,6 +43,7 @@ func init() { cacheVideo = make(map[string]*tele.Video) cacheAlbum = make(map[string]*[]*tele.Photo) + cacheAudio = make(map[string]*tele.Audio) // Open database connection db, err = sql.Open("sqlite", "./data/cache.db") @@ -50,7 +52,7 @@ func init() { return } - execQuery(videoTable, albumTable) + execQuery(videoTable, albumTable, audioTable) load() } diff --git a/utilities.go b/utilities.go index 1401d5b..5a7a2a5 100644 --- a/utilities.go +++ b/utilities.go @@ -71,6 +71,33 @@ func downloadYtDlp(link string) (string, bool) { return filename, hit } +func downloadAudio(link string) (string, bool) { + hit := true + + filename := idGen(link) + ".mp3" + + if _, ok := cacheAudio[filename]; !ok { + // Starts yt-dlp with the arguments to select the best audio + ytDlp := exec.Command("yt-dlp", "-f", "bestaudio", "-f", "mp3", "-q", "-a", "-", "--geo-bypass", "-o", "-") + ytDlp.Stdin = strings.NewReader(link) + out, _ := ytDlp.StdoutPipe() + _ = ytDlp.Start() + + cacheAudio[filename] = &tele.Audio{File: tele.FromReader(out), FileName: filename, MIME: "audio/mp3"} + + go func() { + err := ytDlp.Wait() + if err != nil { + lit.Error(err.Error()) + } + }() + + hit = false + } + + return filename, hit +} + func downloadTikTok(link string) (string, bool, Media) { filename := idGen(link)