YTSFlix_Go/handler/movietorrent.go
2018-11-04 15:58:15 +01:00

146 lines
3.1 KiB
Go

package handler
import (
"fmt"
"github.com/anacrolix/torrent"
"github.com/asticode/go-astisub"
"github.com/odwrtw/yifysubs"
"github.com/sirupsen/logrus"
"gopkg.in/h2non/filetype.v1"
"os"
"path"
"path/filepath"
"strings"
"sync"
"time"
)
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
}