main.js

import { version } from './package.json';
import { CQWebSocket } from 'cq-websocket';
import config from './modules/config';
import saucenao from './modules/saucenao';
import { snDB } from './modules/saucenao';
import whatanime from './modules/whatanime';
import ascii2d from './modules/ascii2d';
import CQ from './modules/CQcode';
import PFCache from './modules/cache';
import Logger from './modules/Logger';
import RandomSeed from 'random-seed';
import sendSetu from './modules/plugin/setu';
import ocr from './modules/plugin/ocr';
import Akhr from './modules/plugin/akhr';
import _ from 'lodash';
import minimist from 'minimist';
import { rmdInit, rmdHandler } from './modules/plugin/reminder';
import broadcast from './modules/broadcast';
import antiBiliMiniApp from './modules/plugin/antiBiliMiniApp';
import logError from './modules/logError';

// 常量
const setting = config.bot;
const rand = RandomSeed.create();
const searchModeOnReg = new RegExp(setting.regs.searchModeOn);
const searchModeOffReg = new RegExp(setting.regs.searchModeOff);

// 初始化
const pfcache = setting.cache.enable ? new PFCache() : null;
if (setting.akhr.enable) Akhr.init().catch(console.error);
if (setting.reminder.enable) rmdInit(replyMsg);

const bot = new CQWebSocket(config.cqws);
const logger = new Logger();

// 好友请求
bot.on('request.friend', context => {
  let approve = setting.autoAddFriend;
  const answers = setting.addFriendAnswers;
  if (approve && answers.length > 0) {
    const comments = context.comment.split('\n');
    try {
      answers.forEach((ans, i) => {
        const a = /(?<=回答:).*/.exec(comments[i * 2 + 1])[0];
        if (ans != a) approve = false;
      });
    } catch (e) {
      console.error(e);
      approve = false;
    }
  }
  if (approve)
    bot('set_friend_add_request', {
      flag: context.flag,
      sub_type: 'invite',
      approve: true,
    });
});

// 加群请求
const groupAddRequests = {};
bot.on('request.group.invite', context => {
  if (setting.autoAddGroup)
    bot('set_group_add_request', {
      flag: context.flag,
      approve: true,
    });
  else groupAddRequests[context.group_id] = context.flag;
});

// 管理员指令
bot.on('message.private', (e, context) => {
  if (context.user_id != setting.admin) return;

  const args = parseArgs(context.message);

  // 允许加群
  const group = args['add-group'];
  if (group && typeof group == 'number') {
    if (typeof groupAddRequests[context.group_id] == 'undefined') {
      replyMsg(context, `将会同意进入群${group}的群邀请`);
      // 注册一次性监听器
      bot.once('request.group.invite', context2 => {
        if (context2.group_id == group) {
          bot('set_group_add_request', {
            flag: context2.flag,
            type: 'invite',
            approve: true,
          });
          replyMsg(context, `已进入群${context2.group_id}`);
          return true;
        }
        return false;
      });
    } else {
      bot('set_group_add_request', {
        flag: groupAddRequests[context.group_id],
        type: 'invite',
        approve: true,
      });
      replyMsg(context, `已进入群${context2.group_id}`);
      delete groupAddRequests[context.group_id];
    }
  }

  if (args.broadcast) broadcast(bot, parseArgs(context.message, false, 'broadcast'));

  // Ban
  const { 'ban-u': bu, 'ban-g': bg } = args;
  if (bu && typeof bu == 'number') {
    Logger.ban('u', bu);
    replyMsg(context, `已封禁用户${bu}`);
  }
  if (bg && typeof bg == 'number') {
    Logger.ban('g', bg);
    replyMsg(context, `已封禁群组${bg}`);
  }

  // 明日方舟
  if (args['update-akhr'])
    Akhr.updateData()
      .then(() => replyMsg(context, '方舟公招数据已更新'))
      .catch(e => {
        logError(e);
        replyMsg(context, '方舟公招数据更新失败，请查看错误日志');
      });

  // 停止程序（利用pm2重启）
  if (args.shutdown) process.exit();
});

// 设置监听器
if (setting.debug) {
  if (setting.enablePM) {
    // 私聊
    bot.on('message.private', debugPrivateAndAtMsg);
  }
  if (setting.enableGM) {
    // 群组@
    bot.on('message.group.@.me', debugPrivateAndAtMsg);
    // 群组
    bot.on('message.group', debugGroupMsg);
  }
} else {
  if (setting.enablePM) {
    // 私聊
    bot.on('message.private', privateAndAtMsg);
  }
  if (setting.enableGM) {
    // 群组@
    bot.on('message.group.@.me', privateAndAtMsg);
    // 群组
    bot.on('message.group', groupMsg);
  }
}

// 连接相关监听
bot
  .on('socket.connecting', (wsType, attempts) => console.log(`${getTime()} 连接中[${wsType}]#${attempts}`))
  .on('socket.failed', (wsType, attempts) => console.log(`${getTime()} 连接失败[${wsType}]#${attempts}`))
  .on('socket.error', (wsType, err) => {
    console.error(`${getTime()} 连接错误[${wsType}]`);
    console.error(err);
  })
  .on('socket.connect', (wsType, sock, attempts) => {
    console.log(`${getTime()} 连接成功[${wsType}]#${attempts}`);
    if (wsType === '/api' && setting.admin > 0) {
      setTimeout(() => {
        bot('send_private_msg', {
          user_id: setting.admin,
          message: `已上线#${attempts}`,
        });
      }, 1000);
    }
  });

// connect
bot.connect();

// 通用处理
function commonHandle(e, context) {
  // 黑名单检测
  if (Logger.checkBan(context.user_id, context.group_id)) return true;

  // 兼容其他机器人
  const startChar = context.message.charAt(0);
  if (startChar == '/' || startChar == '<') return true;

  // 通用指令
  const args = parseArgs(context.message);
  if (args.help) {
    replyMsg(context, 'https://github.com/Tsuk1ko/cq-picsearcher-bot/wiki/%E5%A6%82%E4%BD%95%E9%A3%9F%E7%94%A8');
    return true;
  }
  if (args.version) {
    replyMsg(context, version);
    return true;
  }
  if (args.about) {
    replyMsg(context, 'https://github.com/Tsuk1ko/cq-picsearcher-bot');
    return true;
  }

  // setu
  if (setting.setu.enable) {
    if (sendSetu(context, replyMsg, logger, bot)) return true;
  }

  // reminder
  if (setting.reminder.enable) {
    if (rmdHandler(context)) return true;
  }

  //  反哔哩哔哩小程序
  antiBiliMiniApp(context, replyMsg);

  return false;
}

// 私聊以及群组@的处理
function privateAndAtMsg(e, context) {
  if (commonHandle(e, context)) {
    e.stopPropagation();
    return;
  }

  if (hasImage(context.message)) {
    // 搜图
    e.stopPropagation();
    searchImg(context);
  } else if (context.message.search('--') !== -1) {
    return;
  } else if (!context.group_id && !context.discuss_id) {
    const dbKey = context.message === 'book' ? 'doujin' : context.message;
    const db = snDB[dbKey];
    if (db) {
      logger.smSwitch(0, context.user_id, true);
      logger.smSetDB(0, context.user_id, db);
      return `已临时切换至[${dbKey}]搜图模式√`;
    } else return setting.replys.default;
  } else {
    // 其他指令
    return setting.replys.default;
  }
}

// 调试模式
function debugPrivateAndAtMsg(e, context) {
  if (context.user_id != setting.admin) {
    e.stopPropagation();
    return setting.replys.debug;
  }
  console.log(`${getTime()} 收到私聊消息 qq=${context.user_id}`);
  console.log(_.truncate(context.message, { length: 512, omission: '（字数过多，后续内容不予显示）' }));
  return privateAndAtMsg(e, context);
}

function debugGroupMsg(e, context) {
  if (context.user_id != setting.admin) {
    e.stopPropagation();
    return;
  }
  console.log(`${getTime()} 收到群组消息 group=${context.group_id} qq=${context.user_id}`);
  console.log(_.truncate(context.message, { length: 512, omission: '（字数过多，后续内容不予显示）' }));
  return groupMsg(e, context);
}

// 群组消息处理
function groupMsg(e, context) {
  if (commonHandle(e, context)) {
    e.stopPropagation();
    return;
  }

  // 进入或退出搜图模式
  const { group_id, user_id } = context;

  if (searchModeOnReg.exec(context.message)) {
    // 进入搜图
    e.stopPropagation();
    if (
      logger.smSwitch(group_id, user_id, true, () => {
        replyMsg(context, setting.replys.searchModeTimeout, true);
      })
    )
      replyMsg(context, setting.replys.searchModeOn, true);
    else replyMsg(context, setting.replys.searchModeAlreadyOn, true);
  } else if (searchModeOffReg.exec(context.message)) {
    e.stopPropagation();
    // 退出搜图
    if (logger.smSwitch(group_id, user_id, false)) replyMsg(context, setting.replys.searchModeOff, true);
    else replyMsg(context, setting.replys.searchModeAlreadyOff, true);
  }

  // 搜图模式检测
  let smStatus = logger.smStatus(group_id, user_id);
  if (smStatus) {
    // 获取搜图模式下的搜图参数
    const getDB = () => {
      let cmd = /^(all|pixiv|danbooru|doujin|book|anime)$/.exec(context.message);
      if (cmd) return snDB[cmd[1]] || -1;
      return -1;
    };

    // 切换搜图模式
    const cmdDB = getDB();
    if (cmdDB !== -1) {
      logger.smSetDB(group_id, user_id, cmdDB);
      smStatus = cmdDB;
      replyMsg(context, `已切换至[${context.message}]搜图模式√`);
    }

    // 有图片则搜图
    if (hasImage(context.message)) {
      // 刷新搜图TimeOut
      logger.smSwitch(group_id, user_id, true, () => {
        replyMsg(context, setting.replys.searchModeTimeout, true);
      });
      e.stopPropagation();
      searchImg(context, smStatus);
    }
  } else if (setting.repeat.enable) {
    // 复读（
    // 随机复读，rptLog得到当前复读次数
    if (
      logger.rptLog(group_id, user_id, context.message) >= setting.repeat.times &&
      getRand() <= setting.repeat.probability
    ) {
      logger.rptDone(group_id);
      // 延迟2s后复读
      setTimeout(() => {
        replyMsg(context, context.message);
      }, 2000);
    } else if (getRand() <= setting.repeat.commonProb) {
      // 平时发言下的随机复读
      setTimeout(() => {
        replyMsg(context, context.message);
      }, 2000);
    }
  }
}

/**
 * 搜图
 *
 * @param {object} context
 * @param {number} [customDB=-1]
 * @returns
 */
async function searchImg(context, customDB = -1) {
  const args = parseArgs(context.message);
  const hasWord = word => context.message.indexOf(word) !== -1;

  // OCR
  if (args.ocr) {
    doOCR(context);
    return;
  }

  // 明日方舟
  if (hasWord('akhr') || hasWord('公招')) {
    doAkhr(context);
    return;
  }

  // 决定搜索库
  let db = snDB[setting.saucenaoDefaultDB] || snDB.all;
  if (customDB < 0) {
    if (args.pixiv) db = snDB.pixiv;
    else if (args.danbooru) db = snDB.danbooru;
    else if (args.doujin || args.book) db = snDB.doujin;
    else if (args.anime) db = snDB.anime;
    else if (args.a2d) db = -10001;
    else if (!context.group_id && !context.discuss_id) {
      // 私聊搜图模式
      const sdb = logger.smStatus(0, context.user_id);
      if (sdb) {
        db = sdb;
        logger.smSwitch(0, context.user_id, false);
      }
    }
  } else db = customDB;

  // 得到图片链接并搜图
  const msg = context.message;
  const imgs = getImgs(msg);
  for (const img of imgs) {
    if (args['get-url']) replyMsg(context, img.url.replace(/\/([0-9]+)\/\/\1/, '///').replace(/\?.*$/, ''));
    else {
      // 获取缓存
      let hasCache = false;
      if (setting.cache.enable && !args.purge) {
        const cache = await pfcache.getCache(img.file, db);

        // 如果有缓存
        if (cache) {
          hasCache = true;
          for (const cmsg of cache) {
            replySearchMsgs(context, `${CQ.escape('[缓存]')} ${cmsg}`);
          }
        }
      }

      if (!hasCache) {
        // 检查搜图次数
        if (context.user_id != setting.admin && !logger.canSearch(context.user_id, setting.searchLimit)) {
          replyMsg(context, setting.replys.personLimit, false, true);
          return;
        }

        const needCacheMsgs = [];
        let success = true;
        let useAscii2d = args.a2d;
        let useWhatAnime = args.anime;

        // saucenao
        if (!useAscii2d) {
          const saRet = await saucenao(img.url, db, args.debug || setting.debug);
          if (!saRet.success) success = false;
          if (
            (setting.useAscii2dWhenLowAcc && saRet.lowAcc && (db == snDB.all || db == snDB.pixiv)) ||
            (setting.useAscii2dWhenQuotaExcess && saRet.excess)
          )
            useAscii2d = true;
          if (!saRet.lowAcc && saRet.msg.indexOf('anidb.net') !== -1) useWhatAnime = true;
          if (saRet.msg.length > 0) needCacheMsgs.push(saRet.msg);
          replySearchMsgs(context, saRet.msg, saRet.warnMsg);
        }

        // ascii2d
        if (useAscii2d) {
          const { color, bovw, asErr } = await ascii2d(img.url).catch(asErr => ({
            asErr,
          }));
          if (asErr) {
            const errMsg = (asErr.response && asErr.response.data.length < 50 && `\n${asErr.response.data}`) || '';
            replySearchMsgs(context, `ascii2d 搜索失败${errMsg}`);
            console.error(`${getTime()} [error] ascii2d`);
            logError(asErr);
          } else {
            replySearchMsgs(context, color, bovw);
            needCacheMsgs.push(color);
            needCacheMsgs.push(bovw);
          }
        }

        // 搜番
        if (useWhatAnime) {
          const waRet = await whatanime(img.url, args.debug || setting.debug);
          if (!waRet.success) success = false; //如果搜番有误也视作不成功
          replyMsg(context, waRet.msg, false, true);
          if (waRet.msg.length > 0) needCacheMsgs.push(waRet.msg);
        }

        if (success) logger.doneSearch(context.user_id);

        // 将需要缓存的信息写入数据库
        if (setting.cache.enable && success) {
          await pfcache.addCache(img.file, db, needCacheMsgs);
        }
      }
    }
  }
}

function doOCR(context) {
  const msg = context.message;
  const imgs = getImgs(msg);
  let lang = null;
  const langSearch = /(?<=--lang=)[a-zA-Z]{2,3}/.exec(msg);
  if (langSearch) lang = langSearch[0];

  const handleOcrResult = ret =>
    replyMsg(context, ret.join('\n')).catch(e => {
      replyMsg(context, 'OCR识别发生错误');
      console.error(`${getTime()} [error] OCR`);
      console.error(e);
    });

  for (const img of imgs) {
    ocr.default(img.url, lang).then(handleOcrResult);
  }
}

function doAkhr(context) {
  if (setting.akhr.enable) {
    if (!Akhr.isDataReady()) {
      replyMsg(context, '数据尚未准备完成，请等待一会，或查看日志以检查数据拉取是否出错');
      return;
    }

    const msg = context.message;
    const imgs = getImgs(msg);

    const handleWords = words => {
      //  fix some ...
      if (setting.akhr.ocr == 'ocr.space') words = _.map(words, w => w.replace(/冫口了/g, '治疗'));
      replyMsg(context, CQ.img64(Akhr.getResultImg(words)));
    };

    const handleError = e => {
      replyMsg(context, '词条识别出现错误：\n' + e);
      console.error(`${getTime()} [error] Akhr`);
      console.error(e);
    };

    for (const img of imgs) {
      ocr[setting.akhr.ocr](img.url, 'chs').then(handleWords).catch(handleError);
    }
  } else {
    replyMsg(context, '该功能未开启');
  }
}

/**
 * 从消息中提取图片
 *
 * @param {string} msg
 * @returns 图片URL数组
 */
function getImgs(msg) {
  const reg = /\[CQ:image,file=([^,]+),url=([^\]]+)\]/g;
  const result = [];
  let search = reg.exec(msg);
  while (search) {
    result.push({
      file: search[1],
      url: search[2],
    });
    search = reg.exec(msg);
  }
  return result;
}

/**
 * 判断消息是否有图片
 *
 * @param {string} msg 消息
 * @returns 有则返回true
 */
function hasImage(msg) {
  return msg.indexOf('[CQ:image') !== -1;
}

/**
 * 回复消息
 *
 * @param {object} context 消息对象
 * @param {string} message 回复内容
 * @param {boolean} at 是否at发送者
 */
function replyMsg(context, message, at = false, reply = false) {
  if (typeof message !== 'string' || message.length === 0) return;
  if (context.message_type !== 'private') {
    message = `${reply ? CQ.reply(context.message_id) : ''}${at ? CQ.at(context.user_id) : ''}${message}`;
  }
  const logMsg = setting.debug && _.truncate(message, { length: 512, omission: '（字数过多，后续内容不予显示）' });
  switch (context.message_type) {
    case 'private':
      if (setting.debug) {
        console.log(`${getTime()} 发送私聊消息 qq=${context.user_id}`);
        console.log(logMsg);
      }
      return bot('send_private_msg', {
        user_id: context.user_id,
        message,
      });
    case 'group':
      if (setting.debug) {
        console.log(`${getTime()} 发送群组消息 group=${context.group_id} qq=${context.user_id}`);
        console.log(logMsg);
      }
      return bot('send_group_msg', {
        group_id: context.group_id,
        message,
      });
    case 'discuss':
      if (setting.debug) {
        console.log(`${getTime()} 发送讨论组消息 discuss=${context.discuss_id} qq=${context.user_id}`);
        console.log(logMsg);
      }
      return bot('send_discuss_msg', {
        discuss_id: context.discuss_id,
        message,
      });
  }
}

/**
 * 回复搜图消息
 *
 * @param {object} context 消息对象
 * @param {Array<string>} msgs 回复内容
 */
function replySearchMsgs(context, ...msgs) {
  msgs = msgs.filter(msg => msg && typeof msg === 'string');
  if (msgs.length === 0) return;
  let promises = [];
  //  是否私聊回复
  if (setting.pmSearchResult) {
    switch (context.message_type) {
      case 'group':
      case 'discuss':
        if (!context.pmTipSended) {
          context.pmTipSended = true;
          replyMsg(context, '搜图结果将私聊发送', false, true);
        }
        break;
    }
    promises = msgs.map(msg =>
      bot('send_private_msg', {
        user_id: context.user_id,
        message: msg,
      })
    );
  } else {
    promises = msgs.map(msg => replyMsg(context, msg, false, true));
  }
  return Promise.all(promises);
}

/**
 * 生成随机浮点数
 *
 * @returns 0到100之间的随机浮点数
 */
function getRand() {
  return rand.floatBetween(0, 100);
}

function getTime() {
  return new Date().toLocaleString();
}

function parseArgs(str, enableArray = false, _key = null) {
  const m = minimist(
    str
      .replace(/(--\w+)(?:\s*)(\[CQ:)/g, '$1 $2')
      .replace(/(\[CQ:[^\]]+\])(?:\s*)(--\w+)/g, '$1 $2')
      .split(' '),
    {
      boolean: true,
    }
  );
  if (!enableArray) {
    for (const key in m) {
      if (key == '_') continue;
      if (Array.isArray(m[key])) m[key] = m[key][0];
    }
  }
  if (_key && typeof m[_key] == 'string' && m._.length > 0) m[_key] += ' ' + m._.join(' ');
  return m;
}