anitoru/main.go

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()
	}
}