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 }