Files
cnb/agent/modules/cnb-manager.ts

118 lines
3.3 KiB
TypeScript

import { Result } from '@kevisual/query';
import { CNB } from '../../src/index.ts';
export const getConfig = async (opts: { token?: string }) => {
const res = await fetch('https://kevisual.cn/api/router', {
method: 'POST',
body: JSON.stringify({
path: 'config',
key: 'get',
data: {
key: "cnb_center_config.json"
}
}),
headers: {
'Content-Type': 'application/json',
'Authorization': `Bearer ${opts.token!}`
},
}).then(res => res.json());
return res as Result<{
id: string, key: 'cnb_center_config.json', data: {
CNB_API_KEY: string,
CNB_COOKIE: string
}
}>;
}
type CNBItem = {
username: string,
token: string,
cookie?: string
runAt?: number
owner?: boolean
cnb: CNB
}
export class CNBManager {
cnbMap: Map<string, CNBItem> = new Map()
constructor() {
setInterval(() => {
this.clearExpiredCNB()
}, 1000 * 60 * 30) // 每30分钟清理一次过期的 CNB 实例
}
getDefaultCNB() {
const cnbItem = this.cnbMap.get('default')
if (!cnbItem) {
throw new Error('Default CNB not found')
}
return cnbItem
}
async getCNB(opts?: { username?: string, kevisualToken?: string }): Promise<CNBItem | null> {
if (opts?.username) {
return this.getDefaultCNB()
}
const username = opts?.kevisualToken
const cnbItem = this.cnbMap.get(username)
if (cnbItem) {
cnbItem.runAt = Date.now()
return cnbItem
}
const res = await getConfig({ token: opts?.kevisualToken })
if (res.code === 200) {
const cookie = res.data?.data?.CNB_COOKIE
const token = res.data?.data?.CNB_API_KEY
if (token) {
return this.addCNB({ username, token, cookie })
}
}
return null
}
/**
* 通过上下文获取 CNB 实例(直接返回 cnb 对象)
* @param ctx
* @returns CNB 实例
*/
async getContext(ctx: any) {
const tokenUser = ctx?.state?.tokenUser
const username = tokenUser?.username
if (!username) {
ctx.throw(403, 'Unauthorized')
}
if (username === 'default') {
return this.getDefaultCNB().cnb
}
const kevisualToken = ctx.query?.token;
const item = await this.getCNB({ username, kevisualToken });
if (!item) {
ctx.throw(400, '不存在的 CNB 配置项,请检查 登录 Token 是否正确,或添加 CNB 配置')
}
return item.cnb
}
addCNB(opts: Partial<CNBItem>) {
if (!opts.username || !opts.token) {
throw new Error('username and token are required')
}
const exist = this.cnbMap.get(opts.username)
if (exist) {
exist.runAt = Date.now()
return exist
}
const cnb = opts?.cnb || new CNB({ token: opts.token, cookie: opts.cookie });
opts.cnb = cnb;
opts.runAt = Date.now()
this.cnbMap.set(opts.username, opts as CNBItem)
return opts as CNBItem
}
// 定期清理过期的 CNB 实例,默认过期时间为 1 小时
clearExpiredCNB(expireTime = 1000 * 60 * 60) {
const now = Date.now()
for (const [username, item] of this.cnbMap.entries()) {
if (username === 'default') {
continue
}
if (item.runAt && now - item.runAt > expireTime) {
this.cnbMap.delete(username)
}
}
}
clearUsername(username: string) {
this.cnbMap.delete(username)
}
}