rssmon3/handlers/torrent/main.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"local/args"
	"local/storage"
	"log"
	"net/http"
	"path"
	"regexp"
	"strings"
	"time"

	"github.com/mmcdole/gofeed"
)

const sessionHeader = "X-Transmission-Session-Id"

type Config struct {
	url      string
	vpntor   string
	outdir   string
	interval time.Duration
	last     time.Time
	db       storage.DB
	ctx      context.Context
	can      context.CancelFunc
}

func main() {
	config, err := config()
	if err != nil {
		panic(err)
	}
	log.Println(config)

	for {
		if err := mainLoop(config); err != nil {
			panic(err)
		}
	}
}

func mainLoop(config *Config) error {
	block := config.interval - time.Since(config.last)
	log.Printf("Blocking %v", block)
	select {
	case <-time.After(block):
		if err := pull(config.db, config.vpntor, config.outdir, config.url); err != nil {
			log.Println(err)
		}
		config.last = time.Now()
	case <-config.ctx.Done():
		if err := config.ctx.Err(); err != nil {
			return err
		}
	}
	return nil
}

func config() (*Config, error) {
	as := args.NewArgSet()
	as.Append(args.STRING, "url", "url of rss feed", "http://192.168.0.86:33419/api/tag/torrent")
	as.Append(args.STRING, "vpntor", "url of vpntor", "http://192.168.0.86:9091/transmission/rpc")
	as.Append(args.DURATION, "interval", "interval to check feed", "30m")
	as.Append(args.STRING, "outdir", "save dir", "/data/completed-rss")
	as.Append(args.STRING, "db", "db type", "map")
	as.Append(args.STRING, "addr", "db addr", "")
	as.Append(args.STRING, "user", "db user", "")
	as.Append(args.STRING, "pass", "db pass", "")
	if err := as.Parse(); err != nil {
		return &Config{}, err
	}

	db, err := storage.New(
		storage.TypeFromString(as.Get("db").GetString()),
		as.Get("addr").GetString(),
		as.Get("user").GetString(),
		as.Get("pass").GetString(),
	)
	if err != nil {
		panic(err)
	}

	ctx, can := context.WithCancel(context.Background())
	return &Config{
		url:      as.Get("url").GetString(),
		vpntor:   as.Get("vpntor").GetString(),
		interval: as.Get("interval").GetDuration(),
		outdir:   as.Get("outdir").GetString(),
		db:       db,
		ctx:      ctx,
		can:      can,
	}, nil
}

func pull(db storage.DB, vpntor, outdir, url string) error {
	gfeed, err := getGoFeed(url)
	if err != nil {
		return err
	}
	log.Printf("feed: %v", gfeed.Title)
	for _, item := range gfeed.Items {
		if ok, err := isDone(db, item.Link); err != nil {
			return err
		} else if ok {
			continue
		}
		s, err := getItemContent(item)
		if err != nil {
			return err
		}
		if item.Author == nil {
			item.Author = &gofeed.Person{Name: "."}
		}
		if err := handle(vpntor, path.Join(outdir, item.Author.Name), s); err != nil {
			return err
		}
		if err := db.Set(item.Link, []byte{}); err != nil {
			return err
		}
	}
	return nil
}

func getGoFeed(url string) (*gofeed.Feed, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return gofeed.NewParser().Parse(resp.Body)
}

func getItemContent(item *gofeed.Item) (string, error) {
	s := item.Description
	if s == "" {
		s = item.Content
	}
	if s == "" {
		resp, err := http.Get(item.Link)
		if err != nil {
			return s, err
		}
		defer resp.Body.Close()
		b, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return s, err
		}
		s = string(b)
	}
	return s, nil
}

func isDone(db storage.DB, url string) (bool, error) {
	_, err := db.Get(url)
	if err == storage.ErrNotFound {
		return false, nil
	}
	return true, err
}

func handle(vpntor, outdir, content string) error {
	log.Printf("magnets: %v", findMagnets(content))
	for _, magnet := range findMagnets(content) {
		resp, err := submit(vpntor, outdir, magnet)
		if err != nil {
			return err
		}
		if err := succeeded(resp.Body); err != nil {
			return err
		}
	}
	return nil
}

func findMagnets(s string) []string {
	magnetRegexp := regexp.MustCompile(`magnet:.xt[^ $"]*`)
	return magnetRegexp.FindAllString(s, -1)
}

func submit(vpntor, outdir, magnet string) (*http.Response, error) {
	session, err := getSessionID(vpntor)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest("POST", vpntor, buildReqBody(outdir, magnet))
	if err != nil {
		return nil, err
	}
	req.Header.Add(sessionHeader, session)
	return (&http.Client{}).Do(req)
}

func succeeded(body io.ReadCloser) error {
	defer body.Close()
	b, err := ioutil.ReadAll(body)
	if err != nil {
		return err
	}
	var result struct {
		Result string `json:"result"`
	}
	if err := json.Unmarshal(b, &result); err != nil {
		return err
	}
	if result.Result != "success" {
		return fmt.Errorf("denied: %s", b)
	}
	return nil
}

func buildReqBody(outdir, magnet string) io.Reader {
	return strings.NewReader(fmt.Sprintf(`
      {
         "method": "torrent-add",
         "arguments": {
            "filename": %q,
            "download-dir": %q
         }
      }
   `, magnet, outdir))
}

func getSessionID(vpntor string) (string, error) {
	resp, err := http.Get(vpntor)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	id := resp.Header.Get(sessionHeader)
	if id == "" {
		err = errors.New("session id header not found")
	}
	return id, err
}