bughz/gd-utils
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

update tg.js gd.js by @iwestlin

Browse Source
This commit is contained in:
liaojack8 2020-07-04 01:12:20 +08:00
parent f73cc04880
commit ba07307030
3 changed files with 211 additions and 104 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

134
src/gd.js
View File

@ -12,19 +12,28 @@ const { AUTH, RETRY_LIMIT, PARALLEL_LIMIT, TIMEOUT_BASE, TIMEOUT_MAX, LOG_DELAY,
const { db } = require('../db') const { db } = require('../db')
const { make_table, make_tg_table, make_html, summary } = require('./summary') const { make_table, make_tg_table, make_html, summary } = require('./summary')
const FILE_EXCEED_MSG = '您的团队盘文件数已超限(40万)，停止复制'
const FOLDER_TYPE = 'application/vnd.google-apps.folder' const FOLDER_TYPE = 'application/vnd.google-apps.folder'
const { https_proxy } = process.env const { https_proxy } = process.env
const axins = axios.create(https_proxy ? { httpsAgent: new HttpsProxyAgent(https_proxy) } : {}) const axins = axios.create(https_proxy ? { httpsAgent: new HttpsProxyAgent(https_proxy) } : {})
const SA_BATCH_SIZE = 1000
const SA_FILES = fs.readdirSync(path.join(__dirname, '../sa')).filter(v => v.endsWith('.json')) const SA_FILES = fs.readdirSync(path.join(__dirname, '../sa')).filter(v => v.endsWith('.json'))
SA_FILES.flag = 0
let SA_TOKENS = get_sa_batch()
let SA_TOKENS = SA_FILES.map(filename => { function get_sa_batch () {
const new_flag = SA_FILES.flag + SA_BATCH_SIZE
const files = SA_FILES.slice(SA_FILES.flag, new_flag)
SA_FILES.flag = new_flag
return files.map(filename => {
const gtoken = new GoogleToken({ const gtoken = new GoogleToken({
keyFile: path.join(__dirname, '../sa', filename), keyFile: path.join(__dirname, '../sa', filename),
scope: ['https://www.googleapis.com/auth/drive'] scope: ['https://www.googleapis.com/auth/drive']
}) })
return { gtoken, expires: 0 } return { gtoken, expires: 0 }
}) })
}
handle_exit(() => { handle_exit(() => {
// console.log('handle_exit running') // console.log('handle_exit running')
@ -38,11 +47,13 @@ handle_exit(() => {
async function gen_count_body ({ fid, type, update, service_account }) { async function gen_count_body ({ fid, type, update, service_account }) {
async function update_info () { async function update_info () {
const info = await walk_and_save({ fid, update, service_account }) // 这一步已经将fid记录存入数据库中了 const info = await walk_and_save({ fid, update, service_account }) // 这一步已经将fid记录存入数据库中了
const { summary } = db.prepare('SELECT summary from gd WHERE fid=?').get(fid) const row = db.prepare('SELECT summary from gd WHERE fid=?').get(fid)
return [info, JSON.parse(summary)] if (!row) return []
return [info, JSON.parse(row.summary)]
} }
function render_smy (smy, type) { function render_smy (smy, type) {
if (!smy) return
if (['html', 'curl', 'tg'].includes(type)) { if (['html', 'curl', 'tg'].includes(type)) {
smy = (typeof smy === 'object') ? smy : JSON.parse(smy) smy = (typeof smy === 'object') ? smy : JSON.parse(smy)
const type_func = { const type_func = {
@ -66,7 +77,7 @@ async function gen_count_body ({ fid, type, update, service_account }) {
if (!info) { // 说明上次统计过程中断了 if (!info) { // 说明上次统计过程中断了
[info] = await update_info() [info] = await update_info()
} }
return JSON.stringify(info) return info && JSON.stringify(info)
} }
if (smy) return render_smy(smy, type) if (smy) return render_smy(smy, type)
if (record && record.summary) return render_smy(record.summary, type) if (record && record.summary) return render_smy(record.summary, type)
@ -154,7 +165,7 @@ async function walk_and_save ({ fid, not_teamdrive, update, service_account }) {
const loop = setInterval(() => { const loop = setInterval(() => {
const now = dayjs().format('HH:mm:ss') const now = dayjs().format('HH:mm:ss')
const message = `${now} | 已获取对象 ${result.length} | 排队等候的网络请求 ${limit.pendingCount}` const message = `${now} | 已获取对象 ${result.length} | 网络请求 进行中${limit.activeCount}/排队${limit.pendingCount}`
print_progress(message) print_progress(message)
}, 1000) }, 1000)
@ -248,7 +259,7 @@ async function ls_folder ({ fid, not_teamdrive, service_account }) {
} }
async function gen_headers (use_sa) { async function gen_headers (use_sa) {
use_sa = use_sa && SA_TOKENS.length // use_sa = use_sa && SA_TOKENS.length
const access_token = use_sa ? (await get_sa_token()).access_token : (await get_access_token()) const access_token = use_sa ? (await get_sa_token()).access_token : (await get_access_token())
return { authorization: 'Bearer ' + access_token } return { authorization: 'Bearer ' + access_token }
} }
@ -276,15 +287,17 @@ async function get_access_token () {
return data.access_token return data.access_token
} }
// get_sa_token().catch(console.error)
async function get_sa_token () { async function get_sa_token () {
let tk if (!SA_TOKENS.length) SA_TOKENS = get_sa_batch()
while (SA_TOKENS.length) { while (SA_TOKENS.length) {
tk = get_random_element(SA_TOKENS) const tk = get_random_element(SA_TOKENS)
try { try {
return await real_get_sa_token(tk) return await real_get_sa_token(tk)
} catch (e) { } catch (e) {
console.log(e) console.log(e)
SA_TOKENS = SA_TOKENS.filter(v => v.gtoken !== tk.gtoken) SA_TOKENS = SA_TOKENS.filter(v => v.gtoken !== tk.gtoken)
if (!SA_TOKENS.length) SA_TOKENS = get_sa_batch()
} }
} }
throw new Error('没有可用的SA帐号') throw new Error('没有可用的SA帐号')
@ -302,7 +315,7 @@ function real_get_sa_token (el) {
// console.log('got sa token', tokens) // console.log('got sa token', tokens)
const { access_token, expires_in } = tokens const { access_token, expires_in } = tokens
el.value = access_token el.value = access_token
el.expires = Date.now() + 1000 * expires_in el.expires = Date.now() + 1000 * (expires_in - 60 * 5) // 提前5分钟判定为过期
resolve({ access_token, gtoken }) resolve({ access_token, gtoken })
} }
}) })
@ -323,7 +336,7 @@ function validate_fid (fid) {
return fid.match(reg) return fid.match(reg)
} }
async function create_folder (name, parent, use_sa) { async function create_folder (name, parent, use_sa, limit) {
let url = `https://www.googleapis.com/drive/v3/files` let url = `https://www.googleapis.com/drive/v3/files`
const params = { supportsAllDrives: true } const params = { supportsAllDrives: true }
url += '?' + params_to_query(params) url += '?' + params_to_query(params)
@ -333,18 +346,34 @@ async function create_folder (name, parent, use_sa) {
parents: [parent] parents: [parent]
} }
let retry = 0 let retry = 0
let data let err_message
while (!data && (retry < RETRY_LIMIT)) { while (retry < RETRY_LIMIT) {
try { try {
const headers = await gen_headers(use_sa) const headers = await gen_headers(use_sa)
data = (await axins.post(url, post_data, { headers })).data return (await axins.post(url, post_data, { headers })).data
} catch (err) { } catch (err) {
err_message = err.message
retry++ retry++
handle_error(err) handle_error(err)
const data = err && err.response && err.response.data
const message = data && data.error && data.error.message
if (message && message.toLowerCase().includes('file limit')) {
if (limit) limit.clearQueue()
throw new Error(FILE_EXCEED_MSG)
}
console.log('创建目录重试中：', name, '重试次数：', retry) console.log('创建目录重试中：', name, '重试次数：', retry)
} }
} }
return data throw new Error(err_message + ' 目录名：' + name)
}
async function get_name_by_id (fid) {
try {
const { name } = await get_info_by_id(fid, true)
return name
} catch (e) {
return fid
}
} }
async function get_info_by_id (fid, use_sa) { async function get_info_by_id (fid, use_sa) {
@ -353,7 +382,7 @@ async function get_info_by_id (fid, use_sa) {
includeItemsFromAllDrives: true, includeItemsFromAllDrives: true,
supportsAllDrives: true, supportsAllDrives: true,
corpora: 'allDrives', corpora: 'allDrives',
fields: 'id,name,owners' fields: 'id,name'
} }
url += '?' + params_to_query(params) url += '?' + params_to_query(params)
const headers = await gen_headers(use_sa) const headers = await gen_headers(use_sa)
@ -376,7 +405,7 @@ async function user_choose () {
return answer.value return answer.value
} }
async function copy ({ source, target, name, min_size, update, not_teamdrive, service_account, is_server }) { async function copy ({ source, target, name, min_size, update, not_teamdrive, service_account, dncnr, is_server }) {
target = target || DEFAULT_TARGET target = target || DEFAULT_TARGET
if (!target) throw new Error('目标位置不能为空') if (!target) throw new Error('目标位置不能为空')
@ -384,7 +413,7 @@ async function copy ({ source, target, name, min_size, update, not_teamdrive, se
if (record && record.status === 'copying') return console.log('已有相同源和目的地的任务正在运行，强制退出') if (record && record.status === 'copying') return console.log('已有相同源和目的地的任务正在运行，强制退出')
try { try {
return await real_copy({ source, target, name, min_size, update, not_teamdrive, service_account, is_server }) return await real_copy({ source, target, name, min_size, update, dncnr, not_teamdrive, service_account, is_server })
} catch (err) { } catch (err) {
console.error('复制文件夹出错', err) console.error('复制文件夹出错', err)
const record = db.prepare('select id, status from task where source=? and target=?').get(source, target) const record = db.prepare('select id, status from task where source=? and target=?').get(source, target)
@ -393,8 +422,9 @@ async function copy ({ source, target, name, min_size, update, not_teamdrive, se
} }
// 待解决：如果用户手动ctrl+c中断进程，那么已经发出的请求，就算完成了也不会记录到本地数据库中，所以可能产生重复文件（夹） // 待解决：如果用户手动ctrl+c中断进程，那么已经发出的请求，就算完成了也不会记录到本地数据库中，所以可能产生重复文件（夹）
async function real_copy ({ source, target, name, min_size, update, not_teamdrive, service_account, is_server }) { async function real_copy ({ source, target, name, min_size, update, dncnr, not_teamdrive, service_account, is_server }) {
async function get_new_root () { async function get_new_root () {
if (dncnr) return { id: target }
if (name) { if (name) {
return create_folder(name, target, service_account) return create_folder(name, target, service_account)
} else { } else {
@ -432,12 +462,11 @@ async function real_copy ({ source, target, name, min_size, update, not_teamdriv
root, root,
task_id: record.id task_id: record.id
}) })
await copy_files({ files, mapping: all_mapping, root, task_id: record.id }) await copy_files({ files, service_account, root, mapping: all_mapping, task_id: record.id })
db.prepare('update task set status=?, ftime=? where id=?').run('finished', Date.now(), record.id) db.prepare('update task set status=?, ftime=? where id=?').run('finished', Date.now(), record.id)
return { id: root } return { id: root, task_id: record.id }
} else if (choice === 'restart') { } else if (choice === 'restart') {
const new_root = await get_new_root() const new_root = await get_new_root()
if (!new_root) throw new Error('创建目录失败，请检查您的帐号是否有相应的权限')
const root_mapping = source + ' ' + new_root.id + '\n' const root_mapping = source + ' ' + new_root.id + '\n'
db.prepare('update task set status=?, copied=?, mapping=? where id=?') db.prepare('update task set status=?, copied=?, mapping=? where id=?')
.run('copying', '', root_mapping, record.id) .run('copying', '', root_mapping, record.id)
@ -454,16 +483,15 @@ async function real_copy ({ source, target, name, min_size, update, not_teamdriv
root: new_root.id, root: new_root.id,
task_id: record.id task_id: record.id
}) })
await copy_files({ files, mapping, root: new_root.id, task_id: record.id }) await copy_files({ files, mapping, service_account, root: new_root.id, task_id: record.id })
db.prepare('update task set status=?, ftime=? where id=?').run('finished', Date.now(), record.id) db.prepare('update task set status=?, ftime=? where id=?').run('finished', Date.now(), record.id)
return new_root return { id: new_root.id, task_id: record.id }
} else { } else {
// ctrl+c 退出 // ctrl+c 退出
return console.log('退出程序') return console.log('退出程序')
} }
} else { } else {
const new_root = await get_new_root() const new_root = await get_new_root()
if (!new_root) throw new Error('创建目录失败，请检查您的帐号是否有相应的权限')
const root_mapping = source + ' ' + new_root.id + '\n' const root_mapping = source + ' ' + new_root.id + '\n'
const { lastInsertRowid } = db.prepare('insert into task (source, target, status, mapping, ctime) values (?, ?, ?, ?, ?)').run(source, target, 'copying', root_mapping, Date.now()) const { lastInsertRowid } = db.prepare('insert into task (source, target, status, mapping, ctime) values (?, ?, ?, ?, ?)').run(source, target, 'copying', root_mapping, Date.now())
const arr = await walk_and_save({ fid: source, update, not_teamdrive, service_account }) const arr = await walk_and_save({ fid: source, update, not_teamdrive, service_account })
@ -479,39 +507,33 @@ async function real_copy ({ source, target, name, min_size, update, not_teamdriv
root: new_root.id, root: new_root.id,
task_id: lastInsertRowid task_id: lastInsertRowid
}) })
await copy_files({ files, mapping, root: new_root.id, task_id: lastInsertRowid }) await copy_files({ files, mapping, service_account, root: new_root.id, task_id: lastInsertRowid })
db.prepare('update task set status=?, ftime=? where id=?').run('finished', Date.now(), lastInsertRowid) db.prepare('update task set status=?, ftime=? where id=?').run('finished', Date.now(), lastInsertRowid)
return new_root return { id: new_root.id, task_id: lastInsertRowid }
} }
} }
async function copy_files ({ files, mapping, root, task_id }) { async function copy_files ({ files, mapping, service_account, root, task_id }) {
console.log('\n开始复制文件，总数：', files.length) console.log('\n开始复制文件，总数：', files.length)
const limit = pLimit(PARALLEL_LIMIT) const limit = pLimit(PARALLEL_LIMIT)
let count = 0 let count = 0
const loop = setInterval(() => { const loop = setInterval(() => {
const now = dayjs().format('HH:mm:ss') const now = dayjs().format('HH:mm:ss')
const message = `${now} | 已复制文件数 ${count} | 排队等候的网络请求 ${limit.pendingCount}` const message = `${now} | 已复制文件数 ${count} | 网络请求 进行中${limit.activeCount}/排队${limit.pendingCount}`
print_progress(message) print_progress(message)
}, 1000) }, 1000)
await Promise.all(files.map(async file => { return Promise.all(files.map(async file => {
try {
const { id, parent } = file const { id, parent } = file
const target = mapping[parent] || root const target = mapping[parent] || root
const new_file = await limit(() => copy_file(id, target)) const new_file = await limit(() => copy_file(id, target, service_account, limit))
if (new_file) { if (new_file) {
db.prepare('update task set status=?, copied = copied || ? where id=?')
.run('copying', id + '\n', task_id)
}
count++ count++
} catch (e) { db.prepare('update task set status=?, copied = copied || ? where id=?').run('copying', id + '\n', task_id)
console.error(e)
} }
})) })).finally(() => clearInterval(loop))
clearInterval(loop)
} }
async function copy_file (id, parent) { async function copy_file (id, parent, use_sa, limit) {
let url = `https://www.googleapis.com/drive/v3/files/${id}/copy` let url = `https://www.googleapis.com/drive/v3/files/${id}/copy`
let params = { supportsAllDrives: true } let params = { supportsAllDrives: true }
url += '?' + params_to_query(params) url += '?' + params_to_query(params)
@ -519,7 +541,7 @@ async function copy_file (id, parent) {
let retry = 0 let retry = 0
while (retry < RETRY_LIMIT) { while (retry < RETRY_LIMIT) {
let gtoken let gtoken
if (SA_TOKENS.length) { // 如果有sa文件则优先使用 if (use_sa) {
const temp = await get_sa_token() const temp = await get_sa_token()
gtoken = temp.gtoken gtoken = temp.gtoken
config.headers = { authorization: 'Bearer ' + temp.access_token } config.headers = { authorization: 'Bearer ' + temp.access_token }
@ -534,13 +556,19 @@ async function copy_file (id, parent) {
handle_error(err) handle_error(err)
const data = err && err.response && err.response.data const data = err && err.response && err.response.data
const message = data && data.error && data.error.message const message = data && data.error && data.error.message
if (message && message.toLowerCase().includes('file limit')) {
if (limit) limit.clearQueue()
throw new Error('您的团队盘文件数已超限，停止复制')
}
if (message && message.toLowerCase().includes('rate limit')) { if (message && message.toLowerCase().includes('rate limit')) {
SA_TOKENS = SA_TOKENS.filter(v => v.gtoken !== gtoken) SA_TOKENS = SA_TOKENS.filter(v => v.gtoken !== gtoken)
if (!SA_TOKENS.length) SA_TOKENS = get_sa_batch()
console.log('此帐号触发使用限额，剩余可用service account帐号数量：', SA_TOKENS.length) console.log('此帐号触发使用限额，剩余可用service account帐号数量：', SA_TOKENS.length)
} }
} }
} }
if (!SA_TOKENS.length) { if (!SA_TOKENS.length) {
if (limit) limit.clearQueue()
throw new Error('所有SA帐号流量已用完') throw new Error('所有SA帐号流量已用完')
} else { } else {
console.warn('复制文件失败，文件id: ' + id) console.warn('复制文件失败，文件id: ' + id)
@ -560,7 +588,7 @@ async function create_folders ({ source, old_mapping, folders, root, task_id, se
const loop = setInterval(() => { const loop = setInterval(() => {
const now = dayjs().format('HH:mm:ss') const now = dayjs().format('HH:mm:ss')
const message = `${now} | 已创建目录数 ${count} | 排队等候的网络请求 ${limit.pendingCount}` const message = `${now} | 已创建目录数 ${count} | 网络请求 进行中${limit.activeCount}/排队${limit.pendingCount}`
print_progress(message) print_progress(message)
}, 1000) }, 1000)
@ -569,14 +597,17 @@ async function create_folders ({ source, old_mapping, folders, root, task_id, se
try { try {
const { name, id, parent } = v const { name, id, parent } = v
const target = mapping[parent] || root const target = mapping[parent] || root
const new_folder = await limit(() => create_folder(name, target, service_account)) const new_folder = await limit(() => create_folder(name, target, service_account, limit))
if (!new_folder) throw new Error(name + '创建失败')
count++ count++
mapping[id] = new_folder.id mapping[id] = new_folder.id
const mapping_record = id + ' ' + new_folder.id + '\n' const mapping_record = id + ' ' + new_folder.id + '\n'
db.prepare('update task set status=?, mapping = mapping || ? where id=?').run('copying', mapping_record, task_id) db.prepare('update task set status=?, mapping = mapping || ? where id=?').run('copying', mapping_record, task_id)
} catch (e) { } catch (e) {
console.error('创建目录出错:', v, e) if (e.message === FILE_EXCEED_MSG) {
clearInterval(loop)
throw new Error(FILE_EXCEED_MSG)
}
console.error('创建目录出错:', e.message)
} }
})) }))
folders = folders.filter(v => !mapping[v.id]) folders = folders.filter(v => !mapping[v.id])
@ -627,7 +658,7 @@ async function confirm_dedupe ({ file_number, folder_number }) {
const answer = await prompts({ const answer = await prompts({
type: 'select', type: 'select',
name: 'value', name: 'value',
message: `检测到重复文件${file_number}个，重复目录${folder_number}个，是否删除？`, message: `检测到同位置下重复文件${file_number}个，重复空目录${folder_number}个，是否删除？`,
choices: [ choices: [
{ title: 'Yes', description: '确认删除', value: 'yes' }, { title: 'Yes', description: '确认删除', value: 'yes' },
{ title: 'No', description: '先不删除', value: 'no' } { title: 'No', description: '先不删除', value: 'no' }
@ -637,7 +668,14 @@ async function confirm_dedupe ({ file_number, folder_number }) {
return answer.value return answer.value
} }
// 可以删除文件或文件夹，似乎不会进入回收站 // 将文件或文件夹移入回收站，需要 sa 为 content manager 权限及以上
async function trash_file ({ fid, service_account }) {
const url = `https://www.googleapis.com/drive/v3/files/${fid}?supportsAllDrives=true`
const headers = await gen_headers(service_account)
return axins.patch(url, { trashed: true }, { headers })
}
// 直接删除文件或文件夹，不会进入回收站，需要 sa 为 manager 权限
async function rm_file ({ fid, service_account }) { async function rm_file ({ fid, service_account }) {
const headers = await gen_headers(service_account) const headers = await gen_headers(service_account)
let retry = 0 let retry = 0
@ -677,7 +715,7 @@ async function dedupe ({ fid, update, service_account }) {
let file_count = 0 let file_count = 0
await Promise.all(dupes.map(async v => { await Promise.all(dupes.map(async v => {
try { try {
await limit(() => rm_file({ fid: v.id, service_account })) await limit(() => trash_file({ fid: v.id, service_account }))
if (v.mimeType === FOLDER_TYPE) { if (v.mimeType === FOLDER_TYPE) {
console.log('成功删除文件夹', v.name) console.log('成功删除文件夹', v.name)
folder_count++ folder_count++
@ -710,4 +748,4 @@ function print_progress (msg) {
} }
} }
module.exports = { ls_folder, count, validate_fid, copy, dedupe, copy_file, gen_count_body, real_copy } module.exports = { ls_folder, count, validate_fid, copy, dedupe, copy_file, gen_count_body, real_copy, get_name_by_id }

36
src/router.js
View File

@ -2,21 +2,25 @@ const Router = require('@koa/router')
const { db } = require('../db') const { db } = require('../db')
const { validate_fid, gen_count_body } = require('./gd') const { validate_fid, gen_count_body } = require('./gd')
const { send_count, send_help, send_choice, send_task_info, sm, extract_fid, reply_cb_query, tg_copy, send_all_tasks } = require('./tg') const { send_count, send_help, send_choice, send_task_info, sm, extract_fid, extract_from_text, reply_cb_query, tg_copy, send_all_tasks } = require('./tg')
const { AUTH } = require('../config') const { AUTH, ROUTER_PASSKEY, TG_IPLIST } = require('../config')
const { tg_whitelist } = AUTH const { tg_whitelist } = AUTH
const counting = {} const counting = {}
const router = new Router() const router = new Router()
router.get('/api/gdurl/count', async ctx => { router.get('/api/gdurl/count', async ctx => {
if (!ROUTER_PASSKEY) return ctx.body = 'gd-utils 成功启动'
const { query, headers } = ctx.request const { query, headers } = ctx.request
let { fid, type, update } = query let { fid, type, update, passkey } = query
if (passkey !== ROUTER_PASSKEY) return ctx.body = 'invalid passkey'
if (!validate_fid(fid)) throw new Error('无效的分享ID') if (!validate_fid(fid)) throw new Error('无效的分享ID')
let ua = headers['user-agent'] || '' let ua = headers['user-agent'] || ''
ua = ua.toLowerCase() ua = ua.toLowerCase()
type = (type || '').toLowerCase() type = (type || '').toLowerCase()
// todo type=tree
if (!type) { if (!type) {
if (ua.includes('curl')) { if (ua.includes('curl')) {
type = 'curl' type = 'curl'
@ -38,6 +42,7 @@ router.post('/api/gdurl/tgbot', async ctx => {
const { body } = ctx.request const { body } = ctx.request
console.log('ctx.ip', ctx.ip) // 可以只允许tg服务器的ip console.log('ctx.ip', ctx.ip) // 可以只允许tg服务器的ip
console.log('tg message:', body) console.log('tg message:', body)
if (TG_IPLIST && !TG_IPLIST.includes(ctx.ip)) return ctx.body = 'invalid ip'
ctx.body = '' // 早点释放连接 ctx.body = '' // 早点释放连接
const message = body.message || body.edited_message const message = body.message || body.edited_message
@ -65,23 +70,27 @@ router.post('/api/gdurl/tgbot', async ctx => {
const chat_id = message && message.chat && message.chat.id const chat_id = message && message.chat && message.chat.id
const text = message && message.text && message.text.trim() const text = message && message.text && message.text.trim()
const username = message && message.from && message.from.username let username = message && message.from && message.from.username
if (!chat_id || !text || !tg_whitelist.includes(username)) return console.warn('异常请求') username = username && String(username).toLowerCase()
let user_id = message && message.from && message.from.id
user_id = user_id && String(user_id).toLowerCase()
if (!chat_id || !text || !tg_whitelist.some(v => {
v = String(v).toLowerCase()
return v === username || v === user_id
})) return console.warn('异常请求')
const fid = extract_fid(text) const fid = extract_fid(text) || extract_from_text(text)
const no_fid_commands = ['/task', '/help'] const no_fid_commands = ['/task', '/help']
if (!no_fid_commands.some(cmd => text.startsWith(cmd)) && !validate_fid(fid)) { if (!no_fid_commands.some(cmd => text.startsWith(cmd)) && !validate_fid(fid)) {
return sm({ chat_id, text: '未识别出分享ID' }) return sm({ chat_id, text: '未识别出分享ID' })
} }
if (text.startsWith('/help')) return send_help(chat_id) if (text.startsWith('/help')) return send_help(chat_id)
if (text.startsWith('https://drive.google.com/')) {
return send_choice({ fid, chat_id }).catch(console.error)
}
if (text.startsWith('/count')) { if (text.startsWith('/count')) {
if (counting[fid]) return sm({ chat_id, text: fid + ' 正在统计，请稍等片刻' }) if (counting[fid]) return sm({ chat_id, text: fid + ' 正在统计，请稍等片刻' })
try { try {
counting[fid] = true counting[fid] = true
await send_count({ fid, chat_id }) const update = text.endsWith(' -u')
await send_count({ fid, chat_id, update })
} catch (err) { } catch (err) {
console.error(err) console.error(err)
sm({ chat_id, text: fid + ' 统计失败：' + err.message }) sm({ chat_id, text: fid + ' 统计失败：' + err.message })
@ -89,9 +98,10 @@ router.post('/api/gdurl/tgbot', async ctx => {
delete counting[fid] delete counting[fid]
} }
} else if (text.startsWith('/copy')) { } else if (text.startsWith('/copy')) {
const target = text.replace('/copy', '').trim().split(' ').map(v => v.trim())[1] const target = text.replace('/copy', '').replace(' -u', '').trim().split(' ').map(v => v.trim())[1]
if (target && !validate_fid(target)) return sm({ chat_id, text: `目标ID ${target} 格式不正确` }) if (target && !validate_fid(target)) return sm({ chat_id, text: `目标ID ${target} 格式不正确` })
tg_copy({ fid, target, chat_id }).then(task_id => { const update = text.endsWith(' -u')
tg_copy({ fid, target, chat_id, update }).then(task_id => {
task_id && sm({ chat_id, text: `开始复制，任务ID: ${task_id} 可输入 /task ${task_id} 查询进度` }) task_id && sm({ chat_id, text: `开始复制，任务ID: ${task_id} 可输入 /task ${task_id} 查询进度` })
}) })
} else if (text.startsWith('/task')) { } else if (text.startsWith('/task')) {
@ -106,6 +116,8 @@ router.post('/api/gdurl/tgbot', async ctx => {
return running_tasks.forEach(v => send_task_info({ chat_id, task_id: v.id }).catch(console.error)) return running_tasks.forEach(v => send_task_info({ chat_id, task_id: v.id }).catch(console.error))
} }
send_task_info({ task_id, chat_id }).catch(console.error) send_task_info({ task_id, chat_id }).catch(console.error)
} else if (text.includes('drive.google.com/') || validate_fid(text)) {
return send_choice({ fid: fid || text, chat_id }).catch(console.error)
} else { } else {
sm({ chat_id, text: '暂不支持此命令' }) sm({ chat_id, text: '暂不支持此命令' })
} }

129
src/tg.js
View File

@ -4,15 +4,25 @@ const axios = require('@viegg/axios')
const HttpsProxyAgent = require('https-proxy-agent') const HttpsProxyAgent = require('https-proxy-agent')
const { db } = require('../db') const { db } = require('../db')
const { gen_count_body, validate_fid, real_copy } = require('./gd') const { gen_count_body, validate_fid, real_copy, get_name_by_id } = require('./gd')
const { AUTH, DEFAULT_TARGET } = require('../config') const { AUTH, DEFAULT_TARGET } = require('../config')
const { tg_token } = AUTH const { tg_token } = AUTH
const gen_link = (fid, text) => `<a href="https://drive.google.com/drive/folders/${fid}">${text || fid}</a>`
if (!tg_token) throw new Error('请先在auth.js里设置tg_token') if (!tg_token) throw new Error('请先在auth.js里设置tg_token')
const { https_proxy } = process.env const { https_proxy } = process.env
const axins = axios.create(https_proxy ? { httpsAgent: new HttpsProxyAgent(https_proxy) } : {}) const axins = axios.create(https_proxy ? { httpsAgent: new HttpsProxyAgent(https_proxy) } : {})
module.exports = { send_count, send_help, sm, extract_fid, reply_cb_query, send_choice, send_task_info, send_all_tasks, tg_copy } const FID_TO_NAME = {}
async function get_folder_name (fid) {
let name = FID_TO_NAME[fid]
if (name) return name
name = await get_name_by_id(fid)
return FID_TO_NAME[fid] = name
}
module.exports = { send_count, send_help, sm, extract_fid, reply_cb_query, send_choice, send_task_info, send_all_tasks, tg_copy, extract_from_text }
function send_help (chat_id) { function send_help (chat_id) {
const text = `<pre>[使用帮助] const text = `<pre>[使用帮助]
@ -20,9 +30,9 @@ function send_help (chat_id) {
/help | 返回本条使用说明 /help | 返回本条使用说明
/count shareID | 返回sourceID的文件统计信息, sourceID可以是google drive分享网址本身，也可以是分享ID /count shareID [-u] | 返回sourceID的文件统计信息, sourceID可以是google drive分享网址本身，也可以是分享ID。如果命令最后加上 -u，则无视之前的记录强制从线上获取，适合一段时候后才更新完毕的分享链接。
/copy sourceID targetID | 将sourceID的文件复制到targetID里（会新建一个文件夹），若不填targetID，则会复制到默认位置（在config.js里设置）。返回拷贝任务的taskID /copy sourceID targetID [-u] | 将sourceID的文件复制到targetID里（会新建一个文件夹），若不填targetID，则会复制到默认位置（在config.js里设置）。如果命令最后加上 -u，则无视本地缓存强制从线上获取源文件夹信息。返回拷贝任务的taskID
/task taskID | 返回对应任务的进度信息，若不填则返回所有正在运行的任务进度，若填 all 则返回所有任务列表 /task taskID | 返回对应任务的进度信息，若不填则返回所有正在运行的任务进度，若填 all 则返回所有任务列表
</pre>` </pre>`
@ -60,9 +70,10 @@ async function send_all_tasks (chat_id) {
chat_id, chat_id,
parse_mode: 'HTML', parse_mode: 'HTML',
text: `所有拷贝任务：\n<pre>${text}</pre>` text: `所有拷贝任务：\n<pre>${text}</pre>`
}).catch(async err => { }).catch(err => {
const description = err.response && err.response.data && err.response.data.description // const description = err.response && err.response.data && err.response.data.description
if (description && description.includes('message is too long')) { // if (description && description.includes('message is too long')) {
if (true) {
const text = [headers].concat(records).map(v => v.join('\t')).join('\n') const text = [headers].concat(records).map(v => v.join('\t')).join('\n')
return sm({ chat_id, parse_mode: 'HTML', text: `所有拷贝任务：\n<pre>${text}</pre>` }) return sm({ chat_id, parse_mode: 'HTML', text: `所有拷贝任务：\n<pre>${text}</pre>` })
} }
@ -70,29 +81,53 @@ async function send_all_tasks (chat_id) {
}) })
} }
async function send_task_info ({ task_id, chat_id }) { async function get_task_info (task_id) {
const record = db.prepare('select * from task where id=?').get(task_id) const record = db.prepare('select * from task where id=?').get(task_id)
if (!record) return sm({ chat_id, text: '数据库不存在此任务ID：' + task_id }) if (!record) return {}
const gen_link = fid => `<a href="https://drive.google.com/drive/folders/${fid}">${fid}</a>`
const { source, target, status, copied, mapping, ctime, ftime } = record const { source, target, status, copied, mapping, ctime, ftime } = record
const folder_mapping = mapping && mapping.trim().split('\n')
const new_folder = folder_mapping && folder_mapping[0].split(' ')[1]
const { summary } = db.prepare('select summary from gd where fid=?').get(source) || {} const { summary } = db.prepare('select summary from gd where fid=?').get(source) || {}
const { file_count, folder_count, total_size } = summary ? JSON.parse(summary) : {} const { file_count, folder_count, total_size } = summary ? JSON.parse(summary) : {}
const copied_files = copied ? copied.trim().split('\n').length : 0 const copied_files = copied ? copied.trim().split('\n').length : 0
const copied_folders = mapping ? (mapping.trim().split('\n').length - 1) : 0 const copied_folders = folder_mapping ? (folder_mapping.length - 1) : 0
let text = '任务编号：' + task_id + '\n' let text = '任务编号：' + task_id + '\n'
text += '源ID：' + gen_link(source) + '\n' const folder_name = await get_folder_name(source)
text += '目的ID：' + gen_link(target) + '\n' text += '源文件夹：' + gen_link(source, folder_name) + '\n'
text += '目的位置：' + gen_link(target) + '\n'
text += '新文件夹：' + (new_folder ? gen_link(new_folder) : '暂未创建') + '\n'
text += '任务状态：' + status + '\n' text += '任务状态：' + status + '\n'
text += '创建时间：' + dayjs(ctime).format('YYYY-MM-DD HH:mm:ss') + '\n' text += '创建时间：' + dayjs(ctime).format('YYYY-MM-DD HH:mm:ss') + '\n'
text += '完成时间：' + (ftime ? dayjs(ftime).format('YYYY-MM-DD HH:mm:ss') : '未完成') + '\n' text += '完成时间：' + (ftime ? dayjs(ftime).format('YYYY-MM-DD HH:mm:ss') : '未完成') + '\n'
text += '目录进度：' + copied_folders + '/' + (folder_count === undefined ? '未知数量' : folder_count) + '\n' text += '目录进度：' + copied_folders + '/' + (folder_count === undefined ? '未知数量' : folder_count) + '\n'
text += '文件进度：' + copied_files + '/' + (file_count === undefined ? '未知数量' : file_count) + '\n' text += '文件进度：' + copied_files + '/' + (file_count === undefined ? '未知数量' : file_count) + '\n'
text += '总大小：' + (total_size || '未知大小') text += '合计大小：' + (total_size || '未知大小')
return sm({ chat_id, text, parse_mode: 'HTML' }) const total_count = (folder_count || 0) + (file_count || 0)
return { text, status, total_count }
} }
async function tg_copy ({ fid, target, chat_id }) { // return task_id async function send_task_info ({ task_id, chat_id }) {
const { text, status, total_count } = await get_task_info(task_id)
if (!text) return sm({ chat_id, text: '数据库不存在此任务ID：' + task_id })
const url = `https://api.telegram.org/bot${tg_token}/sendMessage`
let message_id
try {
const { data } = await axins.post(url, { chat_id, text, parse_mode: 'HTML' })
message_id = data && data.result && data.result.message_id
} catch (e) {
console.log('fail to send message to tg', e.message)
}
// get_task_info 在task文件数超大时比较吃cpu，如果超5万就不每10秒更新了
if (!message_id || status !== 'copying' || total_count > 50000) return
const loop = setInterval(async () => {
const url = `https://api.telegram.org/bot${tg_token}/editMessageText`
const { text, status } = await get_task_info(task_id)
if (status !== 'copying') clearInterval(loop)
axins.post(url, { chat_id, message_id, text, parse_mode: 'HTML' }).catch(e => console.error(e.message))
}, 10 * 1000)
}
async function tg_copy ({ fid, target, chat_id, update }) { // return task_id
target = target || DEFAULT_TARGET target = target || DEFAULT_TARGET
if (!target) { if (!target) {
sm({ chat_id, text: '请输入目的地ID或先在config.js里设置默认复制目的地ID(DEFAULT_TARGET)' }) sm({ chat_id, text: '请输入目的地ID或先在config.js里设置默认复制目的地ID(DEFAULT_TARGET)' })
@ -105,17 +140,28 @@ async function tg_copy ({ fid, target, chat_id }) { // return task_id
sm({ chat_id, text: '已有相同源ID和目的ID的任务正在进行，查询进度可输入 /task ' + record.id }) sm({ chat_id, text: '已有相同源ID和目的ID的任务正在进行，查询进度可输入 /task ' + record.id })
return return
} else if (record.status === 'finished') { } else if (record.status === 'finished') {
sm({ chat_id, text: '有相同源ID和目的ID的任务已复制完成，如需重新复制请更换目的地' }) sm({ chat_id, text: `检测到已存在的任务 ${record.id}，开始继续拷贝` })
return
} }
} }
real_copy({ source: fid, target, not_teamdrive: true, service_account: true, is_server: true }) real_copy({ source: fid, update, target, not_teamdrive: true, service_account: true, is_server: true })
.then(folder => { .then(async info => {
if (!record) record = {} // 防止无限循环 if (!record) record = {} // 防止无限循环
if (!folder) return if (!info) return
const link = 'https://drive.google.com/drive/folders/' + folder.id const { task_id } = info
sm({ chat_id, text: `${fid} 复制完成，新文件夹链接：${link}` }) const row = db.prepare('select * from task where id=?').get(task_id)
const { source, target, status, copied, mapping, ctime, ftime } = row
const { summary } = db.prepare('select summary from gd where fid=?').get(source) || {}
const { file_count, folder_count, total_size } = summary ? JSON.parse(summary) : {}
const copied_files = copied ? copied.trim().split('\n').length : 0
const copied_folders = mapping ? (mapping.trim().split('\n').length - 1) : 0
let text = `任务 ${task_id} 复制完成\n`
const name = await get_folder_name(source)
text += '源文件夹：' + gen_link(source, name) + '\n'
text += '目录完成数：' + copied_folders + '/' + folder_count + '\n'
text += '文件完成数：' + copied_files + '/' + file_count + '\n'
sm({ chat_id, text, parse_mode: 'HTML' })
}) })
.catch(err => { .catch(err => {
if (!record) record = {} if (!record) record = {}
@ -145,26 +191,31 @@ function reply_cb_query ({ id, data }) {
}) })
} }
async function send_count ({ fid, chat_id }) { async function send_count ({ fid, chat_id, update }) {
const table = await gen_count_body({ fid, type: 'tg', service_account: true }) const table = await gen_count_body({ fid, update, type: 'tg', service_account: true })
if (!table) return sm({ chat_id, parse_mode: 'HTML', text: gen_link(fid) + ' 信息获取失败' })
const url = `https://api.telegram.org/bot${tg_token}/sendMessage` const url = `https://api.telegram.org/bot${tg_token}/sendMessage`
const gd_link = `https://drive.google.com/drive/folders/${fid}` const gd_link = `https://drive.google.com/drive/folders/${fid}`
const name = await get_folder_name(fid)
return axins.post(url, { return axins.post(url, {
chat_id, chat_id,
parse_mode: 'HTML', parse_mode: 'HTML',
// todo 输出文件名 text: `<pre>源文件夹名称：${name}
text: `<pre>${gd_link} 源链接：${gd_link}
${table}</pre>` ${table}</pre>`
}).catch(async err => { }).catch(async err => {
const description = err.response && err.response.data && err.response.data.description // const description = err.response && err.response.data && err.response.data.description
if (description && description.includes('message is too long')) { // const too_long_msgs = ['request entity too large', 'message is too long']
// if (description && too_long_msgs.some(v => description.toLowerCase().includes(v))) {
if (true) {
const smy = await gen_count_body({ fid, type: 'json', service_account: true }) const smy = await gen_count_body({ fid, type: 'json', service_account: true })
const { file_count, folder_count, total_size } = JSON.parse(smy) const { file_count, folder_count, total_size } = JSON.parse(smy)
return sm({ return sm({
chat_id, chat_id,
parse_mode: 'HTML', parse_mode: 'HTML',
text: `文件统计：<a href="https://drive.google.com/drive/folders/${fid}">${fid}</a>\n<pre> text: `链接：<a href="https://drive.google.com/drive/folders/${fid}">${fid}</a>\n<pre>
表格太长超出telegram消息限制，只显示概要： 表格太长超出telegram消息限制，只显示概要：
目录名称：${name}
文件总数：${file_count} 文件总数：${file_count}
目录总数：${folder_count} 目录总数：${folder_count}
合计大小：${total_size} 合计大小：${total_size}
@ -178,25 +229,31 @@ ${table}</pre>`
function sm (data) { function sm (data) {
const url = `https://api.telegram.org/bot${tg_token}/sendMessage` const url = `https://api.telegram.org/bot${tg_token}/sendMessage`
return axins.post(url, data).catch(err => { return axins.post(url, data).catch(err => {
console.error('fail to post', url, data) // console.error('fail to post', url, data)
console.error(err) console.error('fail to send message to tg:', err.message)
}) })
} }
function extract_fid (text) { function extract_fid (text) {
text = text.replace(/^\/count/, '').replace(/^\/copy/, '').trim() text = text.replace(/^\/count/, '').replace(/^\/copy/, '').replace(/\\/g, '').trim()
const [source, target] = text.split(' ').map(v => v.trim()) const [source, target] = text.split(' ').map(v => v.trim())
if (validate_fid(source)) return source if (validate_fid(source)) return source
try { try {
if (!text.startsWith('http')) text = 'https://' + text if (!text.startsWith('http')) text = 'https://' + text
const u = new URL(text) const u = new URL(text)
if (u.pathname.includes('/folders/')) { if (u.pathname.includes('/folders/')) {
const reg = /\/folders\/([a-zA-Z0-9_-]{10,100})/ const reg = /[^\/?]+$/
const match = u.pathname.match(reg) const match = u.pathname.match(reg)
return match && match[1] return match && match[0]
} }
return u.searchParams.get('id') return u.searchParams.get('id')
} catch (e) { } catch (e) {
return '' return ''
} }
} }
function extract_from_text (text) {
const reg = /https?:\/\/drive.google.com\/[^\s]+/g
const m = text.match(reg)
return m && extract_fid(m[0])
}