refactor, bypass dns filtering using google dns, yts host resolution using qwant

engine/engine.go

@@ -0,0 +1,82 @@
+package engine
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/anacrolix/torrent"
+	"github.com/sirupsen/logrus"
+)
+
+type Engine struct {
+	client      *torrent.Client
+	files       map[string]*MovieTorrent
+	StoragePath string
+}
+
+func NewEngine(storagePath string) (*Engine, error) {
+	logrus.Info("Creating Streaming engine")
+	if storagePath == "" {
+		storagePath = os.TempDir()
+		logrus.Debug("Storage path undefined. Using default temp directory")
+	}
+	logrus.Debugf("Storage Path: %s", storagePath)
+	conf := torrent.NewDefaultClientConfig()
+	conf.DataDir = storagePath
+	c, err := torrent.NewClient(conf)
+	if err != nil {
+		return nil, err
+	}
+	return &Engine{c, map[string]*MovieTorrent{}, storagePath}, nil
+}
+
+func (e *Engine) Close() {
+	e.client.Close()
+}
+
+func (e *Engine) Get(fileName string) (*MovieTorrent, error) {
+	logrus.Infof("Requesting torrent: %s", fileName)
+	if mt, ok := e.files[fileName]; ok {
+		return mt, nil
+	}
+	return nil, fmt.Errorf("not found: %s", fileName)
+}
+
+func (e *Engine) Prepare(magnet string) error {
+	logrus.Debugf("preparing magnet: %s", magnet)
+	t, err := e.client.AddMagnet(magnet)
+	if err != nil {
+		return err
+	}
+	<-t.GotInfo()
+	return nil
+}
+func (e *Engine) Download(magnet string) (*MovieTorrent, error) {
+	logrus.Debugf("downloading magnet: %s", magnet)
+	t, err := e.client.AddMagnet(magnet)
+	if err != nil {
+		return nil, err
+	}
+	<-t.GotInfo()
+	file := getLargestFile(t)
+	file.Download()
+	mt := NewMovieTorrent(file, t)
+	mt.engine = e
+	e.files[mt.FileName] = mt
+	logrus.Infof("Downloading: %s", mt.FileName)
+	return mt, nil
+}
+
+func getLargestFile(t *torrent.Torrent) *torrent.File {
+	var target *torrent.File
+	var maxSize int64
+
+	for _, file := range t.Files() {
+		if maxSize < file.Length() {
+			maxSize = file.Length()
+			target = file
+		}
+	}
+
+	return target
+}

engine/movietorrent.go

@@ -0,0 +1,146 @@
+package engine
+
+import (
+	"fmt"
+	"os"
+	"path"
+	"path/filepath"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/anacrolix/torrent"
+	"github.com/asticode/go-astisub"
+	"github.com/odwrtw/yifysubs"
+	"github.com/sirupsen/logrus"
+	"gopkg.in/h2non/filetype.v1"
+)
+
+type MovieTorrent struct {
+	FileName  string
+	MIME      string
+	Subtitles map[string]string
+	torrent   *torrent.Torrent
+	file      *torrent.File
+	engine    *Engine
+}
+
+func NewMovieTorrent(file *torrent.File, t *torrent.Torrent) *MovieTorrent {
+	fileName := filepath.Base(file.Path())
+	ext := strings.SplitAfter(path.Ext(fileName), ".")[1]
+	MIME := filetype.GetType(ext).MIME.Value
+	return &MovieTorrent{
+		torrent:  t,
+		file:     file,
+		FileName: fileName,
+		MIME:     MIME,
+	}
+}
+
+type Progress struct {
+	Downloaded    uint64
+	Size          uint64
+	DownloadSpeed uint64
+	Percentage    int
+}
+
+func (m *MovieTorrent) NewReader() torrent.Reader {
+	reader := m.file.NewReader()
+	reader.SetReadahead(m.file.Length() / 100)
+	reader.SetResponsive()
+	return reader
+}
+
+func (m *MovieTorrent) Progress() chan Progress {
+	pc := make(chan Progress)
+	var p int64 = 0
+	size := m.torrent.Info().TotalLength()
+	c := m.torrent.Closed()
+	go func() {
+		defer close(pc)
+		for {
+			select {
+			case <-c:
+				logrus.Debugf("%s stopped", m.FileName)
+				return
+			default:
+				if p >= size {
+					logrus.Debugf("%s completed", m.FileName)
+					return
+				}
+				currentProgress := m.torrent.BytesCompleted()
+				percent := int(float64(currentProgress) / float64(size) * 100)
+				pc <- Progress{
+					Downloaded:    uint64(currentProgress),
+					Size:          uint64(size),
+					Percentage:    percent,
+					DownloadSpeed: uint64(currentProgress - p),
+				}
+				p = currentProgress
+				time.Sleep(time.Second)
+
+			}
+		}
+	}()
+	return pc
+}
+
+func (m *MovieTorrent) Cancel() {
+	logrus.Infof("Stopping torrent: %s", m.FileName)
+	m.torrent.Drop()
+}
+
+func (m *MovieTorrent) DownloadSubtitles(imdbID string) error {
+	logrus.Infof("Getting subtitles for: %s", m.FileName)
+	//Create a client
+	client := yifysubs.New("http://yifysubtitles.com")
+
+	//Search subtitles
+	subtitles, err := client.Search(imdbID)
+	if err != nil {
+		return err
+	}
+	m.Subtitles = map[string]string{}
+	basePath := filepath.Join(m.engine.StoragePath, filepath.Dir(m.file.Path()))
+	logrus.Debugf("Subtitles base path: %s", basePath)
+	if _, err := os.Stat(basePath); os.IsNotExist(err) {
+		os.Mkdir(basePath, os.ModePerm)
+	}
+	var wg sync.WaitGroup
+	wg.Add(len(subtitles))
+	mx := sync.RWMutex{}
+	for _, subtitle := range subtitles {
+		go func(sub *yifysubs.Subtitle) {
+			defer wg.Done()
+			name := fmt.Sprintf("%s.vtt", sub.Lang)
+			ref := filepath.Join(filepath.Dir(m.file.Path()), name)
+			p := filepath.Join(basePath, name)
+			if _, err := os.Stat(p); err == nil {
+				mx.Lock()
+				m.Subtitles[name] = ref
+				mx.Unlock()
+				return
+			}
+			file, err := os.Create(p)
+			if err != nil {
+				logrus.Error(err)
+				return
+			}
+			defer file.Close()
+			defer sub.Close()
+			s, _ := astisub.ReadFromSRT(sub)
+
+			if err := s.WriteToWebVTT(file); err != nil {
+				logrus.Debugf("error subtitle: %v", err)
+				os.Remove(p)
+				return
+			}
+			mx.Lock()
+			m.Subtitles[name] = ref
+			mx.Unlock()
+
+		}(subtitle)
+	}
+	wg.Wait()
+	return nil
+}