acfunlive-remake/main.go

// AcFun 直播弹幕录制工具
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/orzogc/acfundanmu"
)

// 命令行参数处理
func argsHandle() {
	const usageStr = "本程序用于录制 AcFun 主播的直播弹幕，在 " + configFile + " 的 uids 里填入要录制弹幕的主播 uid 列表\n" +
		"指定 -danmuburn 时进入弹幕合成模式：读取录制弹幕时生成的 jsonl 文件，合成带头像的弹幕层\n" +
		"弹幕文件也可以直接给 .ass（原版录制的就是这种），会先还原成 jsonl 再合成"

	isHelp := flag.Bool("h", false, "输出本帮助信息")
	configDir = flag.String("config", "", "设置文件所在文件夹，默认是本程序所在文件夹")
	recordDir = flag.String("record", "", "弹幕文件保存的文件夹，默认是本程序所在文件夹")
	danmuburnFile := flag.String("danmuburn", "", "要合成的弹幕文件（.jsonl 或 .ass）路径，指定后进行弹幕合成然后退出，不启动录制")
	burnVideo := flag.String("video", "", "录播视频路径，配合 -danmuburn 把弹幕压进视频；不指定则输出透明头像层 mov")
	burnOut := flag.String("out", "", "合成输出文件路径，默认与弹幕文件同目录")
	burnOffset := flag.Float64("offset", 0, "录播视频相对弹幕录制开始的时间偏移，单位秒，配合 -danmuburn -video 使用")
	burnFFmpeg := flag.String("ffmpeg", "", "ffmpeg 可执行文件路径；不指定时读设置文件里的 ffmpegPath，仍旧找不到再按常见位置查找")
	burnBitrate = *flag.Int("vbitrate", 0, "合成视频的目标码率（Kbps），0 为默认质量，配合 -danmuburn 使用")
	flag.Parse()

	if *isHelp {
		fmt.Println(usageStr)
		flag.PrintDefaults()
		os.Exit(0)
	}
	if flag.NArg() != 0 {
		fmt.Println(usageStr)
		flag.PrintDefaults()
		os.Exit(1)
	}

	if *danmuburnFile != "" {
		loadConfigForDanmuburn()
		// 命令行显式给的优先于设置文件里的
		if *burnFFmpeg != "" {
			config.FFmpegPath = *burnFFmpeg
		}
		runDanmuburn(*danmuburnFile, *burnVideo, *burnOut, *burnOffset)
		os.Exit(0)
	}

	initialize()
}

// 程序初始化
func initialize() {
	exePath, err := os.Executable()
	checkErr(err)
	exeDir = filepath.Dir(exePath)
	if *configDir == "" {
		*configDir = exeDir
	} else {
		info, err := os.Stat(*configDir)
		checkErr(err)
		if !info.IsDir() {
			lPrintErrf("指定的设置文件夹 %s 并不是真正的文件夹", *configDir)
			os.Exit(1)
		}
	}
	if *recordDir == "" {
		*recordDir = exeDir
	} else {
		info, err := os.Stat(*recordDir)
		checkErr(err)
		if !info.IsDir() {
			lPrintErrf("指定的弹幕保存文件夹 %s 并不是真正的文件夹", *recordDir)
			os.Exit(1)
		}
	}
	configFileLocation = filepath.Join(*configDir, configFile)

	if !isConfigFileExist(configFile) {
		data, err := json.MarshalIndent(config, "", "    ")
		checkErr(err)
		err = os.WriteFile(configFileLocation, data, 0644)
		checkErr(err)
		lPrintln("创建设置文件" + configFile + "，请在 uids 里填入要录制弹幕的主播 uid 列表")
	}

	sInfoMap.info = make(map[int]*streamerInfo)
	lInfoMap.info = make(map[string]liveInfo)
	streamers.crt = make(map[int]streamer)
	streamers.old = make(map[int]streamer)

	deviceID, err = acfundanmu.GetDeviceID()
	checkErr(err)

	if ok := fetchAllRooms(); !ok {
		os.Exit(1)
	}
	liveRooms.rooms = liveRooms.newRooms

	loadConfig()
	checkConfig()
	ensureAuthToken()
	buildStreamers()

	for uid, s := range streamers.crt {
		streamers.old[uid] = s
	}
}

func main() {
	argsHandle()

	if len(streamers.crt) == 0 {
		lPrintWarn("请在" + configFile + "的 uids 里填入要录制弹幕的主播 uid 列表")
	} else {
		lPrintln("本程序开始监听主播的直播状态，自动录制直播弹幕")
	}

	if getFFmpeg() == "" {
		lPrintWarn("没有找到 ffmpeg，无法录制直播视频，弹幕合并也无法进行；请安装 ffmpeg 并加入 PATH，" +
			"或把 ffmpeg 放在本程序所在文件夹，或在 web UI 设置区填写「ffmpeg 路径」")
	}
	lPrintf("录播文件格式：%s（可在 web UI 或 %s 的 recordFormat 里改成 flv / ts / mkv）",
		normalizeRecordFormat(config.RecordFormat), configFile)

	mainCh = make(chan controlMsg, 20)

	for _, s := range streamers.crt {
		st := s
		safeGo(st.longID()+"的监听循环", func() { st.cycle("") })
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	mainCtx = ctx

	safeGo("监控设置文件", func() { cycleConfig(ctx) })
	safeGo("获取直播间数据", func() { cycleFetch(ctx) })
	safeGo("清理无用的直播信息", func() { cycleDelKey(ctx) })

	if config.WebPort > 0 {
		safeGo("web UI 服务器", func() { startWebUI(ctx) })
	}

	lPrintln("现在可以输入命令控制本程序，输入 help 查看全部命令的解释")
	safeGo("命令行输入", handleInput)

	// 收到 Ctrl+C（或 Linux 下的终止信号）时走和 quit 命令一样的退出流程，
	// 否则进程会直接消失，ffmpeg 变成孤儿进程继续录制
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
	safeGo("退出信号监听", func() {
		sig := <-sigCh
		lPrintln("收到 " + sig.String() + " 信号，正在退出（再按一次可强制退出）")
		select {
		case mainCh <- controlMsg{c: quit}:
		default:
			return
		}
		<-sigCh
		lPrintln("再次收到退出信号，强制退出")
		os.Exit(1)
	})

Outer:
	for msg := range mainCh {
		switch msg.c {
		case startCycle:
			m := msg
			safeGo(m.s.longID()+"的监听循环", func() { m.s.cycle(m.liveID) })
		case quit:
			// 结束所有 mainCtx 的子 ctx
			cancel()
			// 结束 cycle()
			lPrintln("正在退出各主播的循环")
			sInfoMap.Lock()
			for _, info := range sInfoMap.info {
				// 退出各主播的循环
				if info.ch != nil {
					info.ch <- msg
				}
			}
			sInfoMap.Unlock()
			lInfoMap.RLock()
			for _, info := range lInfoMap.info {
				// 结束录制弹幕
				if info.isDanmu {
					info.danmuCancel()
				}
			}
			lInfoMap.RUnlock()
			// 等待 20 秒，等待弹幕文件写入完成
			time.Sleep(20 * time.Second)
			break Outer
		default:
			lPrintErrf("未知 controlMsg：%+v", msg)
		}
	}

	lPrintln("本程序结束运行")
}