Null-DRM-Official/apps/agent/internal/worker/worker.go

package worker

import (
	"fmt"
	"log"
	"strings"
	"time"

	"drmdecryption/app"
	"drmdecryption/phonecap"

	"drmdecryption/apps/agent/internal/client"
	"drmdecryption/apps/agent/internal/device"
	"drmdecryption/apps/agent/internal/queue"
)

type Worker struct {
	ID          int
	API         *client.Client
	Queue       *queue.Store
	Pool        *device.Pool
	AgentID     string
	CaptureWait time.Duration
	Python      string
	WVD         string
	UserAgent   string
	Log         *log.Logger

	lastDelayLog time.Time
}

func (w *Worker) Loop(stop <-chan struct{}) {
	t := time.NewTicker(2 * time.Second)
	defer t.Stop()
	for {
		select {
		case <-stop:
			return
		case <-t.C:
			w.tick()
		}
	}
}

func (w *Worker) tick() {
	job, ok, err := w.Queue.ClaimNext()
	if err != nil {
		w.logf("claim next: %v", err)
		return
	}
	if !ok {
		if j, until, found, err := w.Queue.NextDelayed(); err == nil && found {
			if time.Since(w.lastDelayLog) > 30*time.Second {
				w.lastDelayLog = time.Now()
				w.logf("no runnable jobs; next is job=%d stream=%s in %s (backoff)",
					j.ID, j.StreamName, until.Round(time.Second))
			}
		}
		return
	}
	w.logf("job=%d running stream=%s app=%s channel=%s reason=%s",
		job.ID, job.StreamName, job.App, job.Channel, job.Reason)
	if err := w.runJob(job); err != nil {
		if device.IsWait(err) || isAuthErr(err) || isUnknownAppErr(err) {
			// Config/environment issues — do not burn attempts or multi-minute backoff.
			w.logf("job=%d waiting: %v", job.ID, err)
			_ = w.Queue.MarkFailed(job.ID, err.Error(), true)
			return
		}
		w.logf("job=%d failed: %v", job.ID, err)
		_ = w.Queue.MarkFailed(job.ID, err.Error(), false)
		return
	}
	_ = w.Queue.MarkOK(job.ID)
	w.logf("job=%d ok stream=%s", job.ID, job.StreamName)
}

func isAuthErr(err error) bool {
	if err == nil {
		return false
	}
	s := strings.ToLower(err.Error())
	return strings.Contains(s, "401") || strings.Contains(s, "unauthorized")
}

func isUnknownAppErr(err error) bool {
	if err == nil {
		return false
	}
	s := strings.ToLower(err.Error())
	return strings.Contains(s, "unknown app")
}

func (w *Worker) runJob(job queue.Job) error {
	a, err := app.Get(job.App)
	if err != nil {
		return err
	}
	label := fmt.Sprintf("job-%d:%s", job.ID, job.StreamName)
	c, serial, err := w.Pool.Acquire(a.Package(), label)
	if err != nil {
		return err
	}
	defer w.Pool.Release(serial)
	_ = w.Queue.SetDevice(job.ID, serial)
	w.logf("job=%d device=%s package=%s", job.ID, serial, a.Package())

	// Short TTL; /ws/agent heartbeats renew while this process is alive.
	if err := w.API.Claim(job.StreamID, w.AgentID, 90); err != nil {
		return fmt.Errorf("streamd claim: %w", err)
	}
	defer func() { _ = w.API.ReleaseClaim(job.StreamID, w.AgentID) }()

	res, err := phonecap.Run(phonecap.Options{
		App:        a,
		Channel:    job.Channel,
		Client:     c,
		Python:     w.Python,
		WVD:        w.WVD,
		UserAgent:  w.UserAgent,
		Wait:       w.CaptureWait,
		CloseApp:   true,
		ClearProxy: true,
	})
	if err != nil {
		return err
	}

	cred := client.Credentials{
		MPD:  res.MPD,
		Key:  res.Key,
		PSSH: res.Session.PSSH,
		Auth: res.Session.Auth,
		PID:  res.Session.PID,
	}
	if err := w.API.PostCredentials(job.StreamID, cred); err != nil {
		return fmt.Errorf("post credentials: %w", err)
	}
	return nil
}

func (w *Worker) logf(format string, args ...any) {
	prefix := fmt.Sprintf("worker/%d: ", w.ID)
	if w.Log != nil {
		w.Log.Printf(prefix+format, args...)
		return
	}
	log.Printf(prefix+format, args...)
}