b站直播弹幕监听

main.go

package main

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/andybalholm/brotli"
	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func main() {

	room := "21721813"
	data, err := danmuID(room)
	if err != nil {
		log.Fatal("danmu: ", err)
	}

	u := url.URL{Scheme: "wss", Host: "tx-sh-live-comet-01.chat.bilibili.com", Path: "/sub"}

	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		log.Fatal("dial: ", err)
	}

	err = c.WriteMessage(websocket.BinaryMessage,
		convertMessage(7, []byte(`{"uid":0,"roomid":`+room+`,"protover":3,"platform":"web","type":2,"key":"`+data["token"].(string)+`"}`)))

	_, _, err = c.ReadMessage()
	if err != nil {
		log.Println("read:", err)
		return
	}

	go func() {
		for {
			//ping
			err = c.WriteMessage(websocket.BinaryMessage,
				convertMessage(2, []byte(`[object Object]`)))
			time.Sleep(time.Second * 30)
		}
	}()

	for {
		_, message, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			return
		}
		msgs := convertRecv(bytes.NewReader(message))
		fmt.Println("msgs: ", len(msgs))
		for _, v := range msgs {
			if strings.Index(string(v.Data), "DANMU_MSG") != -1 {
				log.Printf("recv: %s", v.Data)
			}
		}
	}

}

func convertMessage(msgType int32, data []byte) []byte {
	buf := bytes.NewBuffer([]byte{})
	// length
	binary.Write(buf, binary.BigEndian, int32(16+len(data)))
	//wsBinaryHeaderList
	binary.Write(buf, binary.BigEndian, int16(16))
	binary.Write(buf, binary.BigEndian, int16(1))
	binary.Write(buf, binary.BigEndian, msgType)
	binary.Write(buf, binary.BigEndian, int32(1))
	return append(buf.Bytes(), data...)
}

type msg struct {
	MsgType int
	Data    []byte
}

func convertRecv(r io.Reader) []*msg {
	ret := make([]*msg, 0)

	var length, op, seq int32
	// len
	if err := binary.Read(r, binary.BigEndian, &length); err != nil {
		return ret
	}
	// wsBinaryHeaderList
	var headerLen, ver int16
	binary.Read(r, binary.BigEndian, &headerLen)
	binary.Read(r, binary.BigEndian, &ver)
	binary.Read(r, binary.BigEndian, &op)
	binary.Read(r, binary.BigEndian, &seq)

	for {
		flag := false
		switch ver {
		case 0:
			data := make([]byte, length-16)
			r.Read(data)
			ret = append(ret, &msg{
				MsgType: 0,
				Data:    data,
			})
		case 3:
			br := brotli.NewReader(r)
			ret = append(ret, convertRecv(br)...)
		default:
			flag = true
		}
		if flag {
			break
		}
		// len
		if err := binary.Read(r, binary.BigEndian, &length); err != nil {
			break
		}
		// wsBinaryHeaderList
		var headerLen, ver int16
		binary.Read(r, binary.BigEndian, &headerLen)
		binary.Read(r, binary.BigEndian, &ver)
		binary.Read(r, binary.BigEndian, &op)
		binary.Read(r, binary.BigEndian, &seq)
	}

	return ret
}

func danmuID(roomid string) (map[string]interface{}, error) {
	url := "https://api.live.bilibili.com/xlive/web-room/v1/index/getDanmuInfo?id=" + roomid + "&type=0"
	method := "GET"

	client := &http.Client{
	}
	req, err := http.NewRequest(method, url, nil)

	if err != nil {
		return nil, err
	}
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}

	ret := make(map[string]interface{})

	if err := json.Unmarshal(body, &ret); err != nil {
		return nil, err
	}

	return ret["data"].(map[string]interface{}), nil
}