piaoquan_crawler/common/limit.py

import json
import time
import requests
from datetime import datetime, timedelta
from .aliyun_log import AliyunLogger
from common.scheduling_db import MysqlHelper
from common.db import RedisClient


def generate_expire_time():
    """
    计算出过期时间
    """
    now = datetime.now()

    # 当天晚上12点的时间
    midnight = datetime(now.year, now.month, now.day) + timedelta(days=1)

    # 计算当前时间到当天晚上12点的秒数
    seconds_until_midnight = int((midnight - now).total_seconds())
    return seconds_until_midnight


def bot(name):
    """
    报警机器人
    """
    id_dict = {
        "余海涛": "15869772902",
        "范军": "15200827642",
        "罗情": "15111037095",
        "鲁涛": "18573105114",
        "王雪珂": "13513479926",
        "邓锋": "18175188134",
    }
    url = "https://open.feishu.cn/open-apis/bot/v2/hook/c273ff48-7b7e-4078-b3f7-d69a3d262df9"
    headers = {"Content-Type": "application/json"}
    payload = {
        "msg_type": "interactive",
        "card": {
            "elements": [
                {
                    "tag": "div",
                    "text": {
                        "content": "抓取数量触发限量通知, <at id={}></at>, <at id={}></at>, <at id={}></at>\n".format(
                            id_dict[name], id_dict["邓锋"], id_dict["王雪珂"]
                        ),
                        "tag": "lark_md",
                    },
                },
                {
                    "tag": "div",
                    "text": {
                        "content": "当天已经入库 300 条视频",
                        "tag": "lark_md",
                    },
                },
            ],
            "header": {"title": {"content": "【 通知 】", "tag": "plain_text"}},
        },
    }
    requests.post(url, headers=headers, data=json.dumps(payload))


class AuthorLimit(object):
    """
    账号爬虫限量
    """

    def __init__(self, mode, platform):
        self.mode = mode
        self.platform = platform
        self.limit_tag_dict = {"352": "余海涛", "353": "罗情", "53": "范军", "51": "鲁涛"}

    def find_tag(self, uid):
        """
        通过 uid 去找符合标准的 tag
        """
        sql = f"""select tag from crawler_user_v3 where uid={uid};"""
        result = MysqlHelper.get_values(
            log_type=self.mode, crawler=self.platform, env="prod", sql=sql
        )
        tags = result[0]["tag"]
        if tags:
            tags = tags.split(",")
            for tag in tags:
                if self.limit_tag_dict.get(tag):
                    return tag
        return None

    def author_limitation(self, user_id):
        """
        限制账号， 服务长沙四名同学
        """
        if self.mode == "author":
            tag = self.find_tag(user_id)
            if tag:
                if RedisClient().connect():
                    tag_count = RedisClient().select(tag)
                    if tag_count:
                        tag_count = int(tag_count.decode("utf-8"))
                        if tag_count <= 300:
                            tag_count += 1
                            expire_seconds = generate_expire_time()
                            RedisClient().insert(
                                key=tag, value=tag_count, expire_time=expire_seconds
                            )
                            return True
                        else:
                            # 报警
                            name = self.limit_tag_dict[tag]
                            bot(name)
                            return False
                    else:
                        tag_count = 1
                        expire_seconds = generate_expire_time()
                        RedisClient().insert(
                            key=tag, value=tag_count, expire_time=expire_seconds
                        )
                        return True
                else:
                    return True
            else:
                return True
        else:
            return True