nsfwapp/backend/record_stream_hls.go

// backend\record_stream_hls.go

package main

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"github.com/grafov/m3u8"
)

func getWantedResolutionPlaylistWithHeaders(
	ctx context.Context,
	playlistURL string,
	httpCookie string,
	userAgent string,
	referer string,
) (string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, playlistURL, nil)
	if err != nil {
		return "", err
	}

	if strings.TrimSpace(userAgent) != "" {
		req.Header.Set("User-Agent", strings.TrimSpace(userAgent))
	}
	if strings.TrimSpace(httpCookie) != "" {
		req.Header.Set("Cookie", strings.TrimSpace(httpCookie))
	}
	if strings.TrimSpace(referer) != "" {
		ref := strings.TrimSpace(referer)
		req.Header.Set("Referer", ref)
		if ru, err := url.Parse(ref); err == nil && ru != nil && ru.Scheme != "" && ru.Host != "" {
			req.Header.Set("Origin", ru.Scheme+"://"+ru.Host)
		}
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("HTTP %d beim Abruf der m3u8", resp.StatusCode)
	}

	playlist, listType, err := m3u8.DecodeFrom(resp.Body, true)
	if err != nil {
		return "", fmt.Errorf("m3u8 parse: %w", err)
	}

	if listType == m3u8.MEDIA {
		return playlistURL, nil
	}

	master := playlist.(*m3u8.MasterPlaylist)

	var bestURI string
	var bestHeight int
	var bestFramerate float64

	for _, v := range master.Variants {
		if v == nil || strings.TrimSpace(v.URI) == "" {
			continue
		}

		h := 0
		if v.Resolution != "" {
			parts := strings.Split(v.Resolution, "x")
			if len(parts) == 2 {
				if hh, err := strconv.Atoi(parts[1]); err == nil {
					h = hh
				}
			}
		}

		fr := 30.0
		if v.FrameRate > 0 {
			fr = v.FrameRate
		}

		if h > bestHeight || (h == bestHeight && fr > bestFramerate) {
			bestHeight = h
			bestFramerate = fr
			bestURI = strings.TrimSpace(v.URI)
		}
	}

	if bestURI == "" {
		return "", errors.New("Master-Playlist ohne gültige Varianten")
	}

	baseURL, err := url.Parse(playlistURL)
	if err != nil {
		return "", err
	}
	refURL, err := url.Parse(bestURI)
	if err != nil {
		return "", err
	}

	return baseURL.ResolveReference(refURL).String(), nil
}

func buildFFmpegInputHeaders(httpCookie, userAgent, referer string) string {
	lines := make([]string, 0, 6)

	if strings.TrimSpace(httpCookie) != "" {
		lines = append(lines, fmt.Sprintf("Cookie: %s", strings.TrimSpace(httpCookie)))
	}

	if strings.TrimSpace(referer) != "" {
		ref := strings.TrimSpace(referer)
		lines = append(lines, fmt.Sprintf("Referer: %s", ref))

		if ru, err := url.Parse(ref); err == nil && ru != nil && ru.Scheme != "" && ru.Host != "" {
			lines = append(lines, fmt.Sprintf("Origin: %s://%s", ru.Scheme, ru.Host))
		}
	}

	if strings.TrimSpace(userAgent) != "" {
		lines = append(lines, fmt.Sprintf("User-Agent: %s", strings.TrimSpace(userAgent)))
	}

	if len(lines) == 0 {
		return ""
	}
	return strings.Join(lines, "\r\n") + "\r\n"
}

func handleM3U8Mode(
	ctx context.Context,
	m3u8URL string,
	outFile string,
	job *RecordJob,
	httpCookie string,
	userAgent string,
	referer string,
) error {
	u, err := url.Parse(m3u8URL)
	if err != nil || (u.Scheme != "http" && u.Scheme != "https") {
		return fmt.Errorf("ungültige URL: %q", m3u8URL)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, m3u8URL, nil)
	if err != nil {
		return err
	}
	if strings.TrimSpace(userAgent) != "" {
		req.Header.Set("User-Agent", strings.TrimSpace(userAgent))
	}
	if strings.TrimSpace(httpCookie) != "" {
		req.Header.Set("Cookie", strings.TrimSpace(httpCookie))
	}
	if strings.TrimSpace(referer) != "" {
		req.Header.Set("Referer", strings.TrimSpace(referer))
		if ru, err := url.Parse(strings.TrimSpace(referer)); err == nil && ru != nil && ru.Scheme != "" && ru.Host != "" {
			req.Header.Set("Origin", ru.Scheme+"://"+ru.Host)
		}
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	io.Copy(io.Discard, resp.Body)
	resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP %d beim Abruf der m3u8", resp.StatusCode)
	}

	if strings.TrimSpace(outFile) == "" {
		return errors.New("output file path leer")
	}

	args := []string{
		"-y",
		"-hide_banner",
		"-nostats",
		"-loglevel", "warning",
		"-protocol_whitelist", "file,http,https,tcp,tls,crypto",
		"-http_persistent", "false",
	}

	ua := strings.TrimSpace(userAgent)
	ref := strings.TrimSpace(referer)
	ck := strings.TrimSpace(httpCookie)

	if ua != "" {
		args = append(args, "-user_agent", ua)
	}

	if ref != "" {
		args = append(args, "-referer", ref)
	}

	if ck != "" {
		args = append(args, "-cookies", ck)
	}

	if hdr := buildFFmpegInputHeaders(httpCookie, userAgent, referer); hdr != "" {
		args = append(args, "-headers", hdr)
	}

	args = append(args,
		"-i", m3u8URL,
		"-c", "copy",
		"-movflags", "+faststart",
		outFile,
	)

	cmd := exec.CommandContext(ctx, ffmpegPath, args...)

	var stderr bytes.Buffer
	cmd.Stdout = io.Discard
	cmd.Stderr = &stderr

	stopStat := make(chan struct{})

	if job != nil {
		go func() {
			t := time.NewTicker(1 * time.Second)
			defer t.Stop()

			var last int64
			for {
				select {
				case <-ctx.Done():
					return
				case <-stopStat:
					return
				case <-t.C:
					fi, err := os.Stat(outFile)
					if err != nil || fi == nil || fi.IsDir() {
						continue
					}
					sz := fi.Size()
					if sz > 0 && sz != last {
						jobsMu.Lock()
						job.SizeBytes = sz
						jobsMu.Unlock()
						last = sz
						_ = publishJob(job.ID)
					}
				}
			}
		}()
	}

	err = cmd.Run()
	close(stopStat)

	if err != nil {
		msg := strings.TrimSpace(stderr.String())
		if msg != "" {
			return fmt.Errorf("ffmpeg m3u8 failed: %w: %s", err, msg)
		}
		return fmt.Errorf("ffmpeg m3u8 failed: %w", err)
	}

	return nil
}