146 lines
3.1 KiB
Go
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
|
|
}
|