diff --git a/cli.go b/cli.go old mode 100644 new mode 100755 index 3448b08..6990737 --- a/cli.go +++ b/cli.go @@ -1,211 +1,212 @@ -package main - -import ( - "bufio" - "fmt" - "net" - "os" - "strconv" - "strings" - - e "git.kjao.me/kjao/anitoru/err" -) - -type SubscriptionManager struct { - Subscriptions []string -} - -func NewSubscriptionManager() *SubscriptionManager { - var m SubscriptionManager - m.Get() - return &m -} - -func (m *SubscriptionManager) UpdateSubs(conn net.Conn) { - rawData, err := ReadData(conn) - e.ExitIf(err, "Unable to read from socket.") - m.Subscriptions = strings.Split(string(rawData), "\n") - if m.Subscriptions[0] == "" { - m.Subscriptions = []string{} - } -} - -func (m *SubscriptionManager) Get() { - conn, err := net.Dial("unix", SocketPath) - e.ExitIf(err, "Unable to open socket. Is the daemon running?") - - err = SendData(conn, []byte{'V'}) - e.ExitIf(err, "Unable to write to socket.") - m.UpdateSubs(conn) -} - -func (m *SubscriptionManager) View() { - fmt.Println() - fmt.Println("Current subscriptions:") - for i, sub := range m.Subscriptions { - fmt.Println(fmt.Sprintf("\t[%02d] %v", i+1, sub)) - } - if len(m.Subscriptions) == 0 { - fmt.Println("\t...Empty\n") - return - } - fmt.Println() -} - -func (m *SubscriptionManager) Add() { - fmt.Println() - - // Get available subscriptions. - fmt.Println("Available subscriptions:") - titles, err := GetReleaseSchedule() - if err != nil { - fmt.Println("Unable to get release schedule.") - return - } - - subMap := GetSubMap(&m.Subscriptions) - - var available []string - i := 0 - for _, title := range titles { - _, exists := subMap[title] - if !exists { - i++ - available = append(available, title) - fmt.Println(fmt.Sprintf("\t[%02d] %v", i, title)) - } - } - fmt.Println() - - // Get comma-separated list of additions. - newSubs := []string{} - var add string - var name string - for { - fmt.Print("Add subscriptions. (M for manual): ") - fmt.Scan(&add) - add = strings.ToLower(add) - if add == "m" { - var err error - fmt.Print("Name: ") - in := bufio.NewReader(os.Stdin) - name, err = in.ReadString('\n') - name = name[:len(name)-1] - e.Panic(err) - newSubs = append(newSubs, name) - } else { - addSplit := strings.Split(add, ",") - retry := false - for _, indexString := range addSplit { - index, err := strconv.Atoi(indexString) - if err != nil { - fmt.Println("Error in indices. Try again.") - retry = true - break - } - if index <= 0 || index > len(available) { - fmt.Fprintln(os.Stdout, "Error in indices. %v is invalid. Try again", index) - retry = true - break - } - - newSubs = append(newSubs, available[index-1]) - } - if retry { - continue - } - } - break - } - - var hevc string - fmt.Print("Only HEVC? (Y/n): ") - fmt.Scan(&hevc) - hevc = strings.ToLower(hevc) - var data []byte - if hevc == "n" { - data = []byte{'A', '0'} - } else { - data = []byte{'A', '1'} - } - - // Send addition update to socket - data = append(data, []byte(strings.Join(newSubs, "\n"))...) - conn, err := net.Dial("unix", SocketPath) - e.ExitIf(err, "Unable to open socket. Is the daemon running?") - - err = SendData(conn, data) - e.ExitIf(err, "Unable to write to socket. Is the daemon running?") - m.UpdateSubs(conn) -} - -func (m *SubscriptionManager) Remove() { - fmt.Println() - m.View() - fmt.Println() - - subMap := GetSubMap(&m.Subscriptions) - - // Get comma-separated list of removals. - var remove string - removeSubs := []string{} - for { - fmt.Print("Remove subscriptions: ") - fmt.Scan(&remove) - addSplit := strings.Split(remove, ",") - skip := false - for _, indexString := range addSplit { - index, err := strconv.Atoi(indexString) - if err != nil { - fmt.Println("Error in indices. Try again.") - skip = true - break - } - if index <= 0 || index > len(m.Subscriptions) { - fmt.Fprintln(os.Stdout, "Error in indices. %v is invalid. Try again", index) - skip = true - break - } - removeSubs = append(removeSubs, (m.Subscriptions)[index-1]) - delete(subMap, (m.Subscriptions)[index-1]) - } - if skip { - continue - } - break - } - - // Send removal update to socket - data := []byte{'R'} - data = append(data, []byte(strings.Join(removeSubs, "\n"))...) - conn, err := net.Dial("unix", SocketPath) - e.ExitIf(err, "Unable to open socket. Is the daemon running?") - - err = SendData(conn, data) - e.ExitIf(err, "Unable to write to socket. Is the daemon running?") - m.UpdateSubs(conn) - m.View() -} - -func userCLI() { - manager := NewSubscriptionManager() - - var action string - for { - fmt.Print("View, Add, Remove subscriptions or Quit: ") - fmt.Scan(&action) - action = strings.ToLower(action) - switch action { - case "v": - manager.View() - case "a": - manager.Add() - case "r": - manager.Remove() - case "q": - os.Exit(1) - default: - fmt.Println("Invalid command, try again.") - } - } -} +package main + +import ( + "bufio" + "fmt" + "net" + "os" + "strconv" + "strings" + + e "git.kjao.me/kjao/anitoru/err" +) + +type SubscriptionManager struct { + Subscriptions []string +} + +func NewSubscriptionManager() *SubscriptionManager { + var m SubscriptionManager + m.Get() + return &m +} + +func (m *SubscriptionManager) UpdateSubs(conn net.Conn) { + rawData, err := ReadData(conn) + e.ExitIf(err, "Unable to read from socket.") + m.Subscriptions = strings.Split(string(rawData), "\n") + if m.Subscriptions[0] == "" { + m.Subscriptions = []string{} + } +} + +func (m *SubscriptionManager) Get() { + conn, err := net.Dial("unix", SocketPath) + e.ExitIf(err, "Unable to open socket. Is the daemon running?") + + err = SendData(conn, []byte{'V'}) + e.ExitIf(err, "Unable to write to socket.") + m.UpdateSubs(conn) +} + +func (m *SubscriptionManager) View() { + fmt.Println() + fmt.Println("Current subscriptions:") + for i, sub := range m.Subscriptions { + fmt.Println(fmt.Sprintf("\t[%02d] %v", i+1, sub)) + } + if len(m.Subscriptions) == 0 { + fmt.Println("\t...Empty\n") + return + } + fmt.Println() +} + +func (m *SubscriptionManager) Add() { + fmt.Println() + + // Get available subscriptions. + fmt.Println("Available subscriptions:") + titles, err := GetReleaseSchedule() + if err != nil { + fmt.Println("Unable to get release schedule.") + return + } + + subMap := GetSubMap(&m.Subscriptions) + + var available []string + i := 0 + for _, title := range titles { + title = strings.Split(title, " | ")[0] + _, exists := subMap[title] + if !exists { + i++ + available = append(available, title) + fmt.Println(fmt.Sprintf("\t[%02d] %v", i, title)) + } + } + fmt.Println() + + // Get comma-separated list of additions. + newSubs := []string{} + var add string + var name string + for { + fmt.Print("Add subscriptions. (M for manual): ") + fmt.Scan(&add) + add = strings.ToLower(add) + if add == "m" { + var err error + fmt.Print("Name: ") + in := bufio.NewReader(os.Stdin) + name, err = in.ReadString('\n') + name = name[:len(name)-1] + e.Panic(err) + newSubs = append(newSubs, name) + } else { + addSplit := strings.Split(add, ",") + retry := false + for _, indexString := range addSplit { + index, err := strconv.Atoi(indexString) + if err != nil { + fmt.Println("Error in indices. Try again.") + retry = true + break + } + if index <= 0 || index > len(available) { + fmt.Fprintln(os.Stdout, "Error in indices. %v is invalid. Try again", index) + retry = true + break + } + + newSubs = append(newSubs, available[index-1]) + } + if retry { + continue + } + } + break + } + + var hevc string + fmt.Print("Only HEVC? (Y/n): ") + fmt.Scan(&hevc) + hevc = strings.ToLower(hevc) + var data []byte + if hevc == "n" { + data = []byte{'A', '0'} + } else { + data = []byte{'A', '1'} + } + + // Send addition update to socket + data = append(data, []byte(strings.Join(newSubs, "\n"))...) + conn, err := net.Dial("unix", SocketPath) + e.ExitIf(err, "Unable to open socket. Is the daemon running?") + + err = SendData(conn, data) + e.ExitIf(err, "Unable to write to socket. Is the daemon running?") + m.UpdateSubs(conn) +} + +func (m *SubscriptionManager) Remove() { + fmt.Println() + m.View() + fmt.Println() + + subMap := GetSubMap(&m.Subscriptions) + + // Get comma-separated list of removals. + var remove string + removeSubs := []string{} + for { + fmt.Print("Remove subscriptions: ") + fmt.Scan(&remove) + addSplit := strings.Split(remove, ",") + skip := false + for _, indexString := range addSplit { + index, err := strconv.Atoi(indexString) + if err != nil { + fmt.Println("Error in indices. Try again.") + skip = true + break + } + if index <= 0 || index > len(m.Subscriptions) { + fmt.Fprintln(os.Stdout, "Error in indices. %v is invalid. Try again", index) + skip = true + break + } + removeSubs = append(removeSubs, (m.Subscriptions)[index-1]) + delete(subMap, (m.Subscriptions)[index-1]) + } + if skip { + continue + } + break + } + + // Send removal update to socket + data := []byte{'R'} + data = append(data, []byte(strings.Join(removeSubs, "\n"))...) + conn, err := net.Dial("unix", SocketPath) + e.ExitIf(err, "Unable to open socket. Is the daemon running?") + + err = SendData(conn, data) + e.ExitIf(err, "Unable to write to socket. Is the daemon running?") + m.UpdateSubs(conn) + m.View() +} + +func userCLI() { + manager := NewSubscriptionManager() + + var action string + for { + fmt.Print("View, Add, Remove subscriptions or Quit: ") + fmt.Scan(&action) + action = strings.ToLower(action) + switch action { + case "v": + manager.View() + case "a": + manager.Add() + case "r": + manager.Remove() + case "q": + os.Exit(1) + default: + fmt.Println("Invalid command, try again.") + } + } +} diff --git a/main.go b/main.go old mode 100644 new mode 100755 index 4c39e7e..125de7d --- a/main.go +++ b/main.go @@ -1,593 +1,594 @@ -package main - -import ( - "encoding/binary" - "errors" - "flag" - "fmt" - "log" - "net" - "net/http" - "os" - "os/signal" - "path/filepath" - "strconv" - "strings" - "syscall" - "time" - - e "git.kjao.me/kjao/anitoru/err" - - "github.com/PuerkitoBio/goquery" - "github.com/gookit/config/v2" - "github.com/gookit/config/v2/yamlv3" - "github.com/mmcdole/gofeed" - "github.com/mrobinsn/go-rtorrent/rtorrent" -) - -const ReleaseURL string = "https://www.erai-raws.info/release-schedule/" -const RSSURL string = "https://www.erai-raws.info/feed/?res=1080p&type=magnet&subs%5B0%5D=us&v0=no&d157edc6b50f28b2776442c03d067d56" -const NyaaURL string = "https://nyaa.si/" - -const SocketPath string = "/tmp/anitoru.sock" -const SubscriptionsFile string = "subscriptions.txt" - -var InvalidDotfile error = errors.New("Dotfile invalid.") - -type Configuration struct { - BaseDir string `config:"baseDir"` - RTorrentHost string `config:"rTorrentHost"` - RSSURL string `config:rssURL` - PollingRate int `config:"pollingRate"` - DownloadDir string `config:"downloadDir"` - LogPath string `config:"logPath"` - NotifyError bool `config:"notifyError"` -} - -type Anime struct { - Name string - Episode int - Magnet string - Version int -} - -func NewAnime(link string) Anime { - return Anime{Magnet: link} -} - -func NTFYMessage(msg string) { - http.Post("https://ntfy.kjao.me/anitoru", "text/plain", strings.NewReader(msg)) -} - -func LogNTFY() e.LogOption { - return func(msg string) { - NTFYMessage(msg) - } -} - -func LogStd() e.LogOption { - return func(msg string) { - fmt.Println(msg) - } -} - -func GetSubMap(subscriptions *[]string) map[string]struct{} { - subMap := make(map[string]struct{}) - for _, sub := range *subscriptions { - var d struct{} - subMap[sub] = d - } - return subMap -} - -func GetSubList(subMap *map[string]struct{}) []string { - subscriptions := make([]string, len(*subMap), len(*subMap)) - i := 0 - for sub, _ := range *subMap { - subscriptions[i] = sub - i++ - } - return subscriptions -} - -func (ani *Anime) InfoFromTitle(title string) error { - titleSplit := strings.Split(title, " ") - ani.Name = strings.Join(titleSplit[:len(titleSplit)-2], " ") - - epInfo := strings.Split(titleSplit[len(titleSplit)-1], "v") - ep, err := strconv.Atoi(epInfo[0]) - if err != nil { - err = fmt.Errorf("%w. Full string: %v", err, title) - } - return e.If(err).ThenIf(func() error { - ani.Episode = ep - if len(epInfo) > 1 { - ani.Version, err = strconv.Atoi(epInfo[1]) - return e.If(err).End() - } else { - ani.Version = 1 - return nil - } - }).End() -} - -// func Scheduler(clock string) <-chan int { -// c := make(chan int) -// clockTime, err := time.Parse("05:15", clock) -// FatalErr(err) - -// year, month, day := time.Now().Date() -// scheduleDate := time.Date(year, month, day, clockTime.Hour(), clockTime.Minute(), -// 0, 0, time.UTC).AddDate(0, 0, 1) - -// go func(date time.Time) { -// for { -// if time.Now().After(date) { -// c <- 1 -// date = date.AddDate(0, 0, 1) -// } else { -// time.Sleep(time.Minute) -// } -// } -// }(scheduleDate) -// return c -// } - -func GetReleaseSchedule() ([]string, error) { - var doc *goquery.Document - var titles []string - - res, err := http.Get(ReleaseURL) - defer res.Body.Close() - - return titles, e.If(err).ThenIf(func() error { // The function will run first before return. - doc, err = goquery.NewDocumentFromReader(res.Body) - return err - }).Then(func() { - doc.Find("table .aa_ss_ops_new").Each(func(i int, s *goquery.Selection) { - titles = append(titles, strings.TrimSpace(s.Text())) - }) - }).End() -} - -func GetRSS(rssURL string) ([]Anime, error) { - var animes []Anime - - fp := gofeed.NewParser() - feed, err := fp.ParseURL(rssURL) - - return animes, e.If(err).Then(func() { - for _, item := range feed.Items { - title := strings.Split(item.Title, "[Magnet] ")[1] - if strings.Contains(title, "(HEVC)") { - title = strings.Split(title, " (HEVC)")[0] - } else { - continue - } - - var newVersion bool - if strings.Contains(title, "(AAC 2.0)") { - title = strings.Split(title, " (AAC 2.0)")[0] - } else if strings.Contains(title, "(Repack)") { - newVersion = true - title = strings.Split(title, " (Repack)")[0] - } - anime := NewAnime(item.Link) - if newVersion { - anime.Version += 1 - } - e.LogIf(anime.InfoFromTitle(title)) - animes = append(animes, anime) - } - }).End() -} - -func GetAllAvailable(name string, onlyHEVC bool) ([]Anime, error) { - var animes []Anime - var res *http.Response - var doc *goquery.Document - - req, err := http.NewRequest("GET", NyaaURL, nil) - - return animes, e.If(err).ThenIf(func() error { - q := req.URL.Query() - var query string - if onlyHEVC { - query = "[Erai-raws] " + name + " [HEVC]" - } else { - query = "[Erai-raws] " + name + " [1080p]" - } - query = strings.Replace(query, ":", " ", -1) - query = strings.Replace(query, ".", " ", -1) - q.Add("q", query) - q.Add("s", "id") // Sort by date and descending - q.Add("o", "desc") - req.URL.RawQuery = q.Encode() - - res, err = (&http.Client{}).Do(req) - return err - }).ThenIf(func() error { - defer res.Body.Close() - doc, err = goquery.NewDocumentFromReader(res.Body) - return err - }).Then(func() { - var epMap = make(map[int]struct{}) - doc.Find(".success").Each(func(i int, s *goquery.Selection) { - cols := s.Children() - cols = cols.Next() // Filename - title := cols.Text() - - if !strings.Contains(title, "[ENG]") { - return - } - - isHEVC := strings.Contains(title, "[HEVC]") - if onlyHEVC && !isHEVC { - return - } - - title = strings.Split(title, "[Erai-raws] ")[1] - if strings.Contains(title, "(AAC 2.0)") { - title = strings.Split(title, " (AAC 2.0)")[0] - } else if strings.Contains(title, "(Repack)") { - title = strings.Split(title, " (Repack)")[0] - } else { - is720 := strings.Contains(title, "[720p]") - is1080 := strings.Contains(title, "[1080p]") - - if is720 { - title = strings.Split(title, " [720p]")[0] - } else if is1080 { - title = strings.Split(title, " [1080p]")[0] - } else { - e.Log(fmt.Sprintf("Episode not found, or processed incorrectly. Got: %s.", title), LogNTFY()) - } - - } - - cols = cols.Next() // Download link - link := cols.Children().Next() - anime := NewAnime(link.AttrOr("href", "none")) - e.LogIf(anime.InfoFromTitle(title)) - anime.Name = name // Actual name, since filenames remove special characters. - - _, exists := epMap[anime.Episode] - if !exists { - animes = append(animes, anime) - var s struct{} - epMap[anime.Episode] = s - } - }) - }).End() -} - -func SocketListener(sock net.Listener) chan net.Conn { - c := make(chan net.Conn) - - go func() { - for { - conn, err := sock.Accept() - if e.LogIf(err) { - continue - } - c <- conn - } - }() - - return c -} - -type Daemon struct { - Tor *rtorrent.RTorrent - PollingRate time.Duration - DownloadDir string - RSSURL string - Listener chan net.Conn - Subscriptions map[string]struct{} - LogOptions []e.LogOption - SubscriptionsPath string -} - -func NewDaemon(server string, polling int, download string, socket net.Listener, - notifyError bool, subPath string, rssURL string) Daemon { - // Connect to rTorrent - tor := rtorrent.New(server, false) - name, err := tor.Name() - - e.ExitIf(err, "Unable to connect to rTorrent.") - fmt.Fprintln(os.Stdout, "Connected to", name, "at", server) - - // Create simple socket for subscriptions - listener := SocketListener(socket) - - // Make subscriptions if doesn't exist. - subFile, err := os.OpenFile(subPath, os.O_RDONLY|os.O_CREATE, 0644) - e.ExitIf(err, "Unable to open subscriptions file.") - subFile.Close() - - data, err := os.ReadFile(subPath) - e.ExitIf(err, "Unable to read subscriptions file.") - var subscriptions []string - if len(data) == 0 { - subscriptions = []string{} - } else { - subscriptions = strings.Split(string(data), "\n") - } - subMap := GetSubMap(&subscriptions) - - var opt []e.LogOption - if notifyError { - opt = []e.LogOption{LogNTFY(), LogStd()} - } else { - opt = []e.LogOption{LogStd()} - } - - return Daemon{ - Tor: tor, - PollingRate: time.Duration(polling), - DownloadDir: download, - RSSURL: rssURL, - Listener: listener, - Subscriptions: subMap, - LogOptions: opt, - SubscriptionsPath: subPath, - } -} - -func (d *Daemon) Serve() { - d.CheckRSS() - rssTick := time.Tick(d.PollingRate * time.Minute) - checkTick := time.Tick(10 * time.Second) - for { - select { - case <-rssTick: - e.Log("Checking RSS", LogStd()) - d.CheckRSS() - case <-checkTick: - d.CheckTorrents() - case conn := <-d.Listener: - d.ReadSocket(conn) - conn.Close() - } - } -} - -func (d *Daemon) CheckRSS() { - animes, err := GetRSS(d.RSSURL) - if e.LogIf(err, d.LogOptions...) { - return - } - - for _, anime := range animes { - _, exists := d.Subscriptions[anime.Name] - if !exists { - continue - } - e.LogIf(d.Download(&anime), d.LogOptions...) - } -} - -func (d *Daemon) CheckTorrents() { - torrents, err := d.Tor.GetTorrents(rtorrent.ViewMain) - e.LogIf(err, d.LogOptions...) - for _, torrent := range torrents { - if torrent.Completed { - d.Tor.Delete(torrent) - arr := strings.Split(torrent.Label, "\n") - name := arr[0] - dir := arr[1] - ext := filepath.Ext(torrent.Name) - err := os.Rename(torrent.Path+"/"+torrent.Name, dir+"/"+name+ext) - e.LogIf(err, d.LogOptions...) - e.Log(fmt.Sprintf("Finished downloading %v.", name), LogNTFY()) - } else { - active, err := d.Tor.IsActive(torrent) - e.LogIf(err, d.LogOptions...) - fmt.Println(torrent.Name, "Active:", active) - d.Tor.ResumeTorrent(torrent) - } - } -} - -func (d *Daemon) SendSubs(conn net.Conn) { - subs := strings.Join(GetSubList(&d.Subscriptions), "\n") - err := SendData(conn, []byte(subs)) - e.LogIf(err, d.LogOptions...) -} - -func (d *Daemon) ReadSocket(conn net.Conn) { - rawData, err := ReadData(conn) - if e.LogIf(err, d.LogOptions...) { - return - } - - command := rawData[0] - data := string(rawData[1:len(rawData)]) - e.Log("Socket Command: " + string(command)) - - var write bool - switch command { - case 'A': - onlyHEVC, _ := strconv.ParseBool(string(data[0])) - data = data[1:] - newSubs := strings.Split(data, "\n") - - // Add each subscription - for _, sub := range newSubs { - var s struct{} - d.Subscriptions[sub] = s - } - d.SendSubs(conn) - - for _, sub := range newSubs { - // For each new subscription, get all available animes. - animes, err := GetAllAvailable(sub, onlyHEVC) - if e.LogIf(err, d.LogOptions...) { - continue - } - - for _, anime := range animes { - e.LogIf(d.Download(&anime), d.LogOptions...) - } - } - write = true - case 'R': - for _, sub := range strings.Split(data, "\n") { - delete(d.Subscriptions, sub) - } - write = true - d.SendSubs(conn) - case 'V': - write = false - d.SendSubs(conn) - default: - e.Log("Invalid Command") - } - - if write { // Write subscriptions to file. - f, err := os.Create(d.SubscriptionsPath) - if e.LogIf(err, d.LogOptions...) { - return - } - defer f.Close() - f.Write([]byte(strings.Join(GetSubList(&d.Subscriptions), "\n"))) - f.Sync() - } -} - -func (d *Daemon) Download(ani *Anime) error { - aniDir := filepath.Join(d.DownloadDir, ani.Name) - dotPath := filepath.Join(aniDir, ".anitoru") - var dot map[int]int - - err := os.MkdirAll(aniDir, os.ModePerm) - return e.If(err).ThenIf(func() error { - dot, err = LoadDotfile(dotPath) - return err - }).Then(func() { - vers, exists := dot[ani.Episode] - if exists || vers >= ani.Version { // No need to download. - return - } - name := fmt.Sprintf("%v - %02d", ani.Name, ani.Episode) - d.Tor.Add(ani.Magnet, - rtorrent.DLabel.SetValue(name+"\n"+aniDir), - ) - e.Log(fmt.Sprintf("Queueing download of %v...", name), LogNTFY()) - dot[ani.Episode] = ani.Version - SaveDotfile(dotPath, dot) - }).End() -} - -func SendData(conn net.Conn, data []byte) error { - length := make([]byte, 4) - binary.LittleEndian.PutUint32(length, uint32(len(data))) - _, err := conn.Write(append(length, data...)) - return err -} - -func ReadData(conn net.Conn) ([]byte, error) { - var data []byte - lengthBuf := make([]byte, 4) - _, err := conn.Read(lengthBuf) - return data, e.If(err).Then(func() { - length := binary.LittleEndian.Uint32(lengthBuf) - data = make([]byte, length) - }).ThenIf(func() error { - _, err = conn.Read(data) - return err - }).End() -} - -func LoadDotfile(path string) (map[int]int, error) { - dot := make(map[int]int) - - // Check if exists. - _, err := os.Stat(path) - if errors.Is(err, os.ErrNotExist) { - return dot, nil - } - - data, err := os.ReadFile(path) - return dot, e.If(err).ThenIf(func() error { - eps := strings.Split(string(data), "\n") - for _, epInfo := range eps { - arr := strings.Split(epInfo, " ") - if len(arr) != 2 { - return InvalidDotfile - } - ep, err1 := strconv.Atoi(arr[0]) - v, err2 := strconv.Atoi(arr[1]) - if err1 != nil || err2 != nil { - return InvalidDotfile - } - dot[ep] = v - } - return nil - }).End() -} - -func SaveDotfile(path string, dot map[int]int) error { - var epInfo = make([]string, len(dot), len(dot)) - i := 0 - for ep, v := range dot { - epInfo[i] = fmt.Sprintf("%v %v", ep, v) - i++ - } - return os.WriteFile(path, []byte(strings.Join(epInfo, "\n")), 0666) -} - -func loadConfig(configPath string) Configuration { - config.WithOptions(func(options *config.Options) { - options.DecoderConfig.TagName = "config" - options.ParseEnv = true - options.Readonly = true - }) - config.AddDriver(yamlv3.Driver) - var conf Configuration - e.ExitIf(config.LoadFiles(configPath), "Unable to load config.") - config.Decode(&conf) - return conf -} - -func main() { - daemonize := flag.Bool("daemon", false, "Run the daemon") - config := flag.String("config", "server.yml", "Config path") - flag.Parse() - - if *daemonize { - conf := loadConfig(*config) - - f, err := os.OpenFile(conf.LogPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0744) - e.ExitIf(err, "Unable to open log file.") - defer f.Close() - log.SetOutput(f) - - os.Remove(SocketPath) - sock, err := net.Listen("unix", SocketPath) - e.ExitIf(err, "Unable to open socket.") - defer sock.Close() - err = os.Chmod(SocketPath, 0777) - - // For cleanup, in case of signal. - c := make(chan os.Signal, 1) - signal.Notify(c, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP) - go func() { - for _ = range c { - f.Close() - sock.Close() - os.Exit(1) - } - }() - - // Run daemon - daemon := NewDaemon( - conf.RTorrentHost, conf.PollingRate, conf.DownloadDir, sock, - conf.NotifyError, filepath.Join(conf.BaseDir, SubscriptionsFile), - conf.RSSURL, - ) - daemon.Serve() - } else { - userCLI() - } -} +package main + +import ( + "encoding/binary" + "errors" + "flag" + "fmt" + "log" + "net" + "net/http" + "os" + "os/signal" + "path/filepath" + "strconv" + "strings" + "syscall" + "time" + + e "git.kjao.me/kjao/anitoru/err" + + "github.com/PuerkitoBio/goquery" + "github.com/gookit/config/v2" + "github.com/gookit/config/v2/yamlv3" + "github.com/mmcdole/gofeed" + "github.com/mrobinsn/go-rtorrent/rtorrent" +) + +const ReleaseURL string = "https://www.erai-raws.info/release-schedule/" +const RSSURL string = "https://www.erai-raws.info/feed/?res=1080p&type=magnet&subs%5B0%5D=us&v0=no&d157edc6b50f28b2776442c03d067d56" +const NyaaURL string = "https://nyaa.si/" + +const SocketPath string = "/tmp/anitoru.sock" +const SubscriptionsFile string = "subscriptions.txt" + +var InvalidDotfile error = errors.New("Dotfile invalid.") + +type Configuration struct { + BaseDir string `config:"baseDir"` + RTorrentHost string `config:"rTorrentHost"` + RSSURL string `config:rssURL` + PollingRate int `config:"pollingRate"` + DownloadDir string `config:"downloadDir"` + LogPath string `config:"logPath"` + NotifyError bool `config:"notifyError"` +} + +type Anime struct { + Name string + Episode int + Magnet string + Version int +} + +func NewAnime(link string) Anime { + return Anime{Magnet: link} +} + +func NTFYMessage(msg string) { + http.Post("https://ntfy.kjao.me/anitoru", "text/plain", strings.NewReader(msg)) +} + +func LogNTFY() e.LogOption { + return func(msg string) { + NTFYMessage(msg) + } +} + +func LogStd() e.LogOption { + return func(msg string) { + fmt.Println(msg) + } +} + +func GetSubMap(subscriptions *[]string) map[string]struct{} { + subMap := make(map[string]struct{}) + for _, sub := range *subscriptions { + var d struct{} + subMap[sub] = d + } + return subMap +} + +func GetSubList(subMap *map[string]struct{}) []string { + subscriptions := make([]string, len(*subMap), len(*subMap)) + i := 0 + for sub, _ := range *subMap { + subscriptions[i] = sub + i++ + } + return subscriptions +} + +func (ani *Anime) InfoFromTitle(title string) error { + titleSplit := strings.Split(title, " ") + ani.Name = strings.Join(titleSplit[:len(titleSplit)-2], " ") + ani.Name = strings.Split(ani.Name, " | ")[0] + + epInfo := strings.Split(titleSplit[len(titleSplit)-1], "v") + ep, err := strconv.Atoi(epInfo[0]) + if err != nil { + err = fmt.Errorf("%w. Full string: %v", err, title) + } + return e.If(err).ThenIf(func() error { + ani.Episode = ep + if len(epInfo) > 1 { + ani.Version, err = strconv.Atoi(epInfo[1]) + return e.If(err).End() + } else { + ani.Version = 1 + return nil + } + }).End() +} + +// func Scheduler(clock string) <-chan int { +// c := make(chan int) +// clockTime, err := time.Parse("05:15", clock) +// FatalErr(err) + +// year, month, day := time.Now().Date() +// scheduleDate := time.Date(year, month, day, clockTime.Hour(), clockTime.Minute(), +// 0, 0, time.UTC).AddDate(0, 0, 1) + +// go func(date time.Time) { +// for { +// if time.Now().After(date) { +// c <- 1 +// date = date.AddDate(0, 0, 1) +// } else { +// time.Sleep(time.Minute) +// } +// } +// }(scheduleDate) +// return c +// } + +func GetReleaseSchedule() ([]string, error) { + var doc *goquery.Document + var titles []string + + res, err := http.Get(ReleaseURL) + defer res.Body.Close() + + return titles, e.If(err).ThenIf(func() error { // The function will run first before return. + doc, err = goquery.NewDocumentFromReader(res.Body) + return err + }).Then(func() { + doc.Find("table .aa_ss_ops_new").Each(func(i int, s *goquery.Selection) { + titles = append(titles, strings.TrimSpace(s.Text())) + }) + }).End() +} + +func GetRSS(rssURL string) ([]Anime, error) { + var animes []Anime + + fp := gofeed.NewParser() + feed, err := fp.ParseURL(rssURL) + + return animes, e.If(err).Then(func() { + for _, item := range feed.Items { + title := strings.Split(item.Title, "[Magnet] ")[1] + if strings.Contains(title, "(HEVC)") { + title = strings.Split(title, " (HEVC)")[0] + } else { + continue + } + + var newVersion bool + if strings.Contains(title, "(AAC 2.0)") { + title = strings.Split(title, " (AAC 2.0)")[0] + } else if strings.Contains(title, "(Repack)") { + newVersion = true + title = strings.Split(title, " (Repack)")[0] + } + anime := NewAnime(item.Link) + if newVersion { + anime.Version += 1 + } + e.LogIf(anime.InfoFromTitle(title)) + animes = append(animes, anime) + } + }).End() +} + +func GetAllAvailable(name string, onlyHEVC bool) ([]Anime, error) { + var animes []Anime + var res *http.Response + var doc *goquery.Document + + req, err := http.NewRequest("GET", NyaaURL, nil) + + return animes, e.If(err).ThenIf(func() error { + q := req.URL.Query() + var query string + if onlyHEVC { + query = "[Erai-raws] " + name + " [HEVC]" + } else { + query = "[Erai-raws] " + name + " [1080p]" + } + query = strings.Replace(query, ":", " ", -1) + query = strings.Replace(query, ".", " ", -1) + q.Add("q", query) + q.Add("s", "id") // Sort by date and descending + q.Add("o", "desc") + req.URL.RawQuery = q.Encode() + + res, err = (&http.Client{}).Do(req) + return err + }).ThenIf(func() error { + defer res.Body.Close() + doc, err = goquery.NewDocumentFromReader(res.Body) + return err + }).Then(func() { + var epMap = make(map[int]struct{}) + doc.Find(".success").Each(func(i int, s *goquery.Selection) { + cols := s.Children() + cols = cols.Next() // Filename + title := cols.Text() + + if !strings.Contains(title, "[ENG]") { + return + } + + isHEVC := strings.Contains(title, "[HEVC]") + if onlyHEVC && !isHEVC { + return + } + + title = strings.Split(title, "[Erai-raws] ")[1] + if strings.Contains(title, "(AAC 2.0)") { + title = strings.Split(title, " (AAC 2.0)")[0] + } else if strings.Contains(title, "(Repack)") { + title = strings.Split(title, " (Repack)")[0] + } else { + is720 := strings.Contains(title, "[720p]") + is1080 := strings.Contains(title, "[1080p]") + + if is720 { + title = strings.Split(title, " [720p]")[0] + } else if is1080 { + title = strings.Split(title, " [1080p]")[0] + } else { + e.Log(fmt.Sprintf("Episode not found, or processed incorrectly. Got: %s.", title), LogNTFY()) + } + + } + + cols = cols.Next() // Download link + link := cols.Children().Next() + anime := NewAnime(link.AttrOr("href", "none")) + e.LogIf(anime.InfoFromTitle(title)) + anime.Name = name // Actual name, since filenames remove special characters. + + _, exists := epMap[anime.Episode] + if !exists { + animes = append(animes, anime) + var s struct{} + epMap[anime.Episode] = s + } + }) + }).End() +} + +func SocketListener(sock net.Listener) chan net.Conn { + c := make(chan net.Conn) + + go func() { + for { + conn, err := sock.Accept() + if e.LogIf(err) { + continue + } + c <- conn + } + }() + + return c +} + +type Daemon struct { + Tor *rtorrent.RTorrent + PollingRate time.Duration + DownloadDir string + RSSURL string + Listener chan net.Conn + Subscriptions map[string]struct{} + LogOptions []e.LogOption + SubscriptionsPath string +} + +func NewDaemon(server string, polling int, download string, socket net.Listener, + notifyError bool, subPath string, rssURL string) Daemon { + // Connect to rTorrent + tor := rtorrent.New(server, false) + name, err := tor.Name() + + e.ExitIf(err, "Unable to connect to rTorrent.") + fmt.Fprintln(os.Stdout, "Connected to", name, "at", server) + + // Create simple socket for subscriptions + listener := SocketListener(socket) + + // Make subscriptions if doesn't exist. + subFile, err := os.OpenFile(subPath, os.O_RDONLY|os.O_CREATE, 0644) + e.ExitIf(err, "Unable to open subscriptions file.") + subFile.Close() + + data, err := os.ReadFile(subPath) + e.ExitIf(err, "Unable to read subscriptions file.") + var subscriptions []string + if len(data) == 0 { + subscriptions = []string{} + } else { + subscriptions = strings.Split(string(data), "\n") + } + subMap := GetSubMap(&subscriptions) + + var opt []e.LogOption + if notifyError { + opt = []e.LogOption{LogNTFY(), LogStd()} + } else { + opt = []e.LogOption{LogStd()} + } + + return Daemon{ + Tor: tor, + PollingRate: time.Duration(polling), + DownloadDir: download, + RSSURL: rssURL, + Listener: listener, + Subscriptions: subMap, + LogOptions: opt, + SubscriptionsPath: subPath, + } +} + +func (d *Daemon) Serve() { + d.CheckRSS() + rssTick := time.Tick(d.PollingRate * time.Minute) + checkTick := time.Tick(10 * time.Second) + for { + select { + case <-rssTick: + e.Log("Checking RSS", LogStd()) + d.CheckRSS() + case <-checkTick: + d.CheckTorrents() + case conn := <-d.Listener: + d.ReadSocket(conn) + conn.Close() + } + } +} + +func (d *Daemon) CheckRSS() { + animes, err := GetRSS(d.RSSURL) + if e.LogIf(err, d.LogOptions...) { + return + } + + for _, anime := range animes { + _, exists := d.Subscriptions[anime.Name] + if !exists { + continue + } + e.LogIf(d.Download(&anime), d.LogOptions...) + } +} + +func (d *Daemon) CheckTorrents() { + torrents, err := d.Tor.GetTorrents(rtorrent.ViewMain) + e.LogIf(err, d.LogOptions...) + for _, torrent := range torrents { + if torrent.Completed { + d.Tor.Delete(torrent) + arr := strings.Split(torrent.Label, "\n") + name := arr[0] + dir := arr[1] + ext := filepath.Ext(torrent.Name) + err := os.Rename(torrent.Path+"/"+torrent.Name, dir+"/"+name+ext) + e.LogIf(err, d.LogOptions...) + e.Log(fmt.Sprintf("Finished downloading %v.", name), LogNTFY()) + } else { + active, err := d.Tor.IsActive(torrent) + e.LogIf(err, d.LogOptions...) + fmt.Println(torrent.Name, "Active:", active) + d.Tor.ResumeTorrent(torrent) + } + } +} + +func (d *Daemon) SendSubs(conn net.Conn) { + subs := strings.Join(GetSubList(&d.Subscriptions), "\n") + err := SendData(conn, []byte(subs)) + e.LogIf(err, d.LogOptions...) +} + +func (d *Daemon) ReadSocket(conn net.Conn) { + rawData, err := ReadData(conn) + if e.LogIf(err, d.LogOptions...) { + return + } + + command := rawData[0] + data := string(rawData[1:len(rawData)]) + e.Log("Socket Command: " + string(command)) + + var write bool + switch command { + case 'A': + onlyHEVC, _ := strconv.ParseBool(string(data[0])) + data = data[1:] + newSubs := strings.Split(data, "\n") + + // Add each subscription + for _, sub := range newSubs { + var s struct{} + d.Subscriptions[sub] = s + } + d.SendSubs(conn) + + for _, sub := range newSubs { + // For each new subscription, get all available animes. + animes, err := GetAllAvailable(sub, onlyHEVC) + if e.LogIf(err, d.LogOptions...) { + continue + } + + for _, anime := range animes { + e.LogIf(d.Download(&anime), d.LogOptions...) + } + } + write = true + case 'R': + for _, sub := range strings.Split(data, "\n") { + delete(d.Subscriptions, sub) + } + write = true + d.SendSubs(conn) + case 'V': + write = false + d.SendSubs(conn) + default: + e.Log("Invalid Command") + } + + if write { // Write subscriptions to file. + f, err := os.Create(d.SubscriptionsPath) + if e.LogIf(err, d.LogOptions...) { + return + } + defer f.Close() + f.Write([]byte(strings.Join(GetSubList(&d.Subscriptions), "\n"))) + f.Sync() + } +} + +func (d *Daemon) Download(ani *Anime) error { + aniDir := filepath.Join(d.DownloadDir, ani.Name) + dotPath := filepath.Join(aniDir, ".anitoru") + var dot map[int]int + + err := os.MkdirAll(aniDir, os.ModePerm) + return e.If(err).ThenIf(func() error { + dot, err = LoadDotfile(dotPath) + return err + }).Then(func() { + vers, exists := dot[ani.Episode] + if exists || vers >= ani.Version { // No need to download. + return + } + name := fmt.Sprintf("%v - %02d", ani.Name, ani.Episode) + d.Tor.Add(ani.Magnet, + rtorrent.DLabel.SetValue(name+"\n"+aniDir), + ) + e.Log(fmt.Sprintf("Queueing download of %v...", name), LogNTFY()) + dot[ani.Episode] = ani.Version + SaveDotfile(dotPath, dot) + }).End() +} + +func SendData(conn net.Conn, data []byte) error { + length := make([]byte, 4) + binary.LittleEndian.PutUint32(length, uint32(len(data))) + _, err := conn.Write(append(length, data...)) + return err +} + +func ReadData(conn net.Conn) ([]byte, error) { + var data []byte + lengthBuf := make([]byte, 4) + _, err := conn.Read(lengthBuf) + return data, e.If(err).Then(func() { + length := binary.LittleEndian.Uint32(lengthBuf) + data = make([]byte, length) + }).ThenIf(func() error { + _, err = conn.Read(data) + return err + }).End() +} + +func LoadDotfile(path string) (map[int]int, error) { + dot := make(map[int]int) + + // Check if exists. + _, err := os.Stat(path) + if errors.Is(err, os.ErrNotExist) { + return dot, nil + } + + data, err := os.ReadFile(path) + return dot, e.If(err).ThenIf(func() error { + eps := strings.Split(string(data), "\n") + for _, epInfo := range eps { + arr := strings.Split(epInfo, " ") + if len(arr) != 2 { + return InvalidDotfile + } + ep, err1 := strconv.Atoi(arr[0]) + v, err2 := strconv.Atoi(arr[1]) + if err1 != nil || err2 != nil { + return InvalidDotfile + } + dot[ep] = v + } + return nil + }).End() +} + +func SaveDotfile(path string, dot map[int]int) error { + var epInfo = make([]string, len(dot), len(dot)) + i := 0 + for ep, v := range dot { + epInfo[i] = fmt.Sprintf("%v %v", ep, v) + i++ + } + return os.WriteFile(path, []byte(strings.Join(epInfo, "\n")), 0666) +} + +func loadConfig(configPath string) Configuration { + config.WithOptions(func(options *config.Options) { + options.DecoderConfig.TagName = "config" + options.ParseEnv = true + options.Readonly = true + }) + config.AddDriver(yamlv3.Driver) + var conf Configuration + e.ExitIf(config.LoadFiles(configPath), "Unable to load config.") + config.Decode(&conf) + return conf +} + +func main() { + daemonize := flag.Bool("daemon", false, "Run the daemon") + config := flag.String("config", "server.yml", "Config path") + flag.Parse() + + if *daemonize { + conf := loadConfig(*config) + + f, err := os.OpenFile(conf.LogPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0744) + e.ExitIf(err, "Unable to open log file.") + defer f.Close() + log.SetOutput(f) + + os.Remove(SocketPath) + sock, err := net.Listen("unix", SocketPath) + e.ExitIf(err, "Unable to open socket.") + defer sock.Close() + err = os.Chmod(SocketPath, 0777) + + // For cleanup, in case of signal. + c := make(chan os.Signal, 1) + signal.Notify(c, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP) + go func() { + for _ = range c { + f.Close() + sock.Close() + os.Exit(1) + } + }() + + // Run daemon + daemon := NewDaemon( + conf.RTorrentHost, conf.PollingRate, conf.DownloadDir, sock, + conf.NotifyError, filepath.Join(conf.BaseDir, SubscriptionsFile), + conf.RSSURL, + ) + daemon.Serve() + } else { + userCLI() + } +}