send_msg.py

import requests
import json
import os
import re
import logging
from time import sleep
from typing import List, Dict

# Configure logging
logging.basicConfig(
    level=logging.INFO, 
    format='%(asctime)s - %(levelname)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

class MessageSender:
    def __init__(self):
        # 储存各URL的发送计数
        self.usage_count: Dict[str, int] = {"gd": 0, "gx": 0, "gz":0, "pb": 0}
        
        # 初始化URL字典
        self.url_dict: Dict[str, List[str]] = {
            "gd": [],
            "gx": [],
            "gz": [],
            "pb": []
        }
        
        # 加载环境变量中的URL
        self._load_urls()
    
    def _load_urls(self):
        """从环境变量加载URL列表"""
        try:
            if 'gk_gd' in os.environ:
                self.url_dict["gd"] = re.split("@|#|\n", os.environ.get("gk_gd",''))
            
            if 'gk_gx' in os.environ:
                self.url_dict["gx"] = re.split("@|#|\n", os.environ.get("gk_gx",''))
            
            if 'gk_gz' in os.environ:
                self.url_dict["gz"] = re.split("@|#|\n", os.environ.get("gk_gz",''))

            if 'db_pb' in os.environ:
                self.url_dict["pb"] = re.split("@|#|\n", os.environ.get("db_pb",''))
        except Exception as e:
            logger.error(f"Error loading URLs from environment: {e}")
    
    def _send_single_message(self, url: str, msg: str) -> bool:
        """发送单条消息"""
        headers = {
            'Content-Type': 'application/json'
        }
        
        data = {
            "msgtype": "text",
            "text": {
                "content": msg
            }
        }
        
        try:
            response = requests.post(url, headers=headers, data=json.dumps(data), timeout=10)
            response.raise_for_status()  # 对于错误的HTTP状态码抛出异常
            
            logger.info(f"Message sent successfully to {url}: {response.status_code}")
            return True
        except requests.RequestException as e:
            logger.error(f"Failed to send message to {url}: {e}")
            return False
    
    def send_messages(self, place: str, messages: List[str]):
        """
        发送消息列表
        
        :param place: 地区标识 ('gd' 或 'gx','gz','pb')
        :param messages: 要发送的消息列表
        """
        # 验证输入
        if place not in ["gd", "gx","gz", "pb"]:
            logger.error(f"Invalid place: {place}. Must be 'gd', 'gx', 'gz', or 'pb'.")
            return
        
        # 为每个目标place维护独立的索引
        current_indexes = {"gd": 0, "pb": 0}

        for msg in messages:
            # 根据消息内容决定发送目标
            if "【广东公考】" in msg:
                if "【编制岗位】" in msg or "【编制状态不确定】" in msg:
                    target_place = "gd"
                elif "【非编制岗位】" in msg:
                    target_place = "pb"
            else:
                # 如果没有编制标识，使用原始的place参数
                target_place = place
            
            urls = self.url_dict.get(target_place, [])
            
            if not urls:
                logger.warning(f"No URLs found for target place: {target_place}")
                continue
            
            current_index = current_indexes.get(target_place, 0)
            
            while True:
                # 获取当前URL
                url = urls[current_index]
                
                # 发送消息
                if self._send_single_message(url, msg):
                    # 增加计数
                    self.usage_count[target_place] += 1
                    
                    # 检查是否到达上限
                    if self.usage_count[target_place] >= 20:
                        # 重置计数并切换到下一个地址
                        self.usage_count[target_place] = 0
                        current_index += 1
                        
                        # 检查是否还有下一个地址
                        if current_index >= len(urls):
                            logger.info(f"All URLs for {target_place} have reached their limit. Pausing for a minute.")
                            sleep(60)  # 等待一分钟
                            current_index = 0  # 从第一个地址重新开始
                    
                    # 更新索引
                    current_indexes[target_place] = current_index
                    
                    # 暂停以避免速率限制
                    sleep(0.5)
                    break
                else:
                    # 如果发送失败，切换到下一个URL
                    current_index += 1
                    if current_index >= len(urls):
                        logger.error(f"All URLs for {target_place} failed. Cannot send message.")
                        break
                    
                    # 更新索引
                    current_indexes[target_place] = current_index