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"` 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]) 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() ([]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] title = strings.Split(title, " [1080p]")[0] if title[len(title)-6:] != "(HEVC)" { continue } title = title[:len(title)-7] anime := NewAnime(item.Link) e.LogIf(anime.InfoFromTitle(title)) animes = append(animes, anime) } }).End() } func GetAllAvailable(name string) ([]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() query := "[Erai-raws] " + name + " [HEVC]" 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, "[HEVC]") || !strings.Contains(title, "[ENG]") { return } title = strings.Split(title, "[Erai-raws] ")[1] if strings.Contains(title, "(AAC 2.0)") { title = strings.Split(title, " (AAC 2.0)")[0] } else { title = strings.Split(title, " [1080p]")[0] } 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 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) 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, Listener: listener, Subscriptions: subMap, LogOptions: opt, SubscriptionsPath: subPath, } } func (d *Daemon) Serve() { //d.CheckRSS() for { select { case <-time.After(d.PollingRate * time.Minute): d.CheckRSS() case <-time.After(10 * time.Second): d.CheckTorrents() case conn := <-d.Listener: d.ReadSocket(conn) conn.Close() } } } func (d *Daemon) CheckRSS() { animes, err := GetRSS() 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)) newSubs := strings.Split(data, "\n") var write bool switch command { case 'A': // 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) 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), ) daemon.Serve() } else { userCLI() } }