Rename to hkt.sh

projects/news-bot/storage.py

@@ -0,0 +1,187 @@
+"""
+数据存储模块 - 使用 JSON 文件持久化
+"""
+import json
+import os
+import time
+import logging
+from pathlib import Path
+from threading import Lock
+
+logger = logging.getLogger(__name__)
+
+DATA_DIR = os.environ.get("DATA_DIR", "/opt/news-bot/data")
+
+# 文件锁，防止并发写入
+_locks = {}
+
+
+def _get_lock(name: str) -> Lock:
+    if name not in _locks:
+        _locks[name] = Lock()
+    return _locks[name]
+
+
+def _ensure_dir():
+    """确保数据目录存在"""
+    Path(DATA_DIR).mkdir(parents=True, exist_ok=True)
+
+
+def _read_json(filename: str, default=None):
+    """读取 JSON 文件"""
+    filepath = os.path.join(DATA_DIR, filename)
+    try:
+        if os.path.exists(filepath):
+            with open(filepath, "r", encoding="utf-8") as f:
+                return json.load(f)
+    except Exception as e:
+        logger.error(f"读取 {filename} 失败: {e}")
+    return default if default is not None else {}
+
+
+def _write_json(filename: str, data):
+    """写入 JSON 文件"""
+    _ensure_dir()
+    filepath = os.path.join(DATA_DIR, filename)
+    try:
+        with open(filepath, "w", encoding="utf-8") as f:
+            json.dump(data, f, ensure_ascii=False, indent=2)
+    except Exception as e:
+        logger.error(f"写入 {filename} 失败: {e}")
+
+
+# ========== 新闻存储 ==========
+
+def save_news(news_list: list):
+    """保存新闻列表（追加，自动去重）"""
+    with _get_lock("news"):
+        existing = _read_json("news.json", [])
+        existing_ids = {n.get("id") for n in existing}
+        added = 0
+        for item in news_list:
+            if item.get("id") not in existing_ids:
+                existing.append(item)
+                existing_ids.add(item["id"])
+                added += 1
+        # 只保留最近 2000 条
+        if len(existing) > 2000:
+            existing = existing[-2000:]
+        _write_json("news.json", existing)
+        return added
+
+
+def get_news(limit=50, since_ts=0) -> list:
+    """获取新闻，可按时间过滤"""
+    news = _read_json("news.json", [])
+    if since_ts:
+        news = [n for n in news if n.get("timestamp", 0) >= since_ts]
+    return news[-limit:]
+
+
+def get_news_by_score(min_score=6, since_ts=0) -> list:
+    """按评分获取新闻"""
+    news = get_news(limit=500, since_ts=since_ts)
+    return [n for n in news if n.get("score", 0) >= min_score]
+
+
+# ========== 已推送记录 ==========
+
+def mark_pushed(news_ids: list):
+    """标记新闻为已推送"""
+    with _get_lock("pushed"):
+        pushed = set(_read_json("pushed.json", []))
+        pushed.update(news_ids)
+        # 只保留最近 5000 条
+        pushed_list = list(pushed)[-5000:]
+        _write_json("pushed.json", pushed_list)
+
+
+def is_pushed(news_id: str) -> bool:
+    """检查新闻是否已推送"""
+    pushed = set(_read_json("pushed.json", []))
+    return news_id in pushed
+
+
+def get_pushed_ids() -> set:
+    """获取所有已推送 ID"""
+    return set(_read_json("pushed.json", []))
+
+
+# ========== 用户设置 ==========
+
+DEFAULT_SETTINGS = {
+    "push_enabled": True,
+    "batch_interval": 30,  # 汇总推送间隔（分钟）
+    "min_score": 6,  # 最低推送评分
+    "instant_score": 8,  # 即时推送评分阈值
+    "sources": {
+        "jin10": True,
+        "wallstreet": True,
+        "kr36": True,
+        "sina": True,
+    },
+    "keywords": [
+        "AI", "英伟达", "NVIDIA", "AMD", "存储", "光模块", "算力",
+        "华为", "腾讯", "阿里", "字节", "小米",
+        "收购", "突破", "泄露", "内幕", "传闻", "重磅",
+        "A股", "港股", "美股",
+    ],
+}
+
+
+def get_settings() -> dict:
+    """获取用户设置"""
+    settings = _read_json("settings.json", None)
+    if settings is None:
+        settings = DEFAULT_SETTINGS.copy()
+        _write_json("settings.json", settings)
+    # 补全缺失字段
+    for k, v in DEFAULT_SETTINGS.items():
+        if k not in settings:
+            settings[k] = v
+    return settings
+
+
+def update_settings(updates: dict):
+    """更新用户设置"""
+    with _get_lock("settings"):
+        settings = get_settings()
+        settings.update(updates)
+        _write_json("settings.json", settings)
+    return settings
+
+
+def toggle_source(source_name: str) -> bool:
+    """切换信息源开关，返回新状态"""
+    with _get_lock("settings"):
+        settings = get_settings()
+        current = settings.get("sources", {}).get(source_name, True)
+        settings.setdefault("sources", {})[source_name] = not current
+        _write_json("settings.json", settings)
+        return not current
+
+
+def add_keyword(keyword: str) -> bool:
+    """添加关键词，返回是否成功"""
+    with _get_lock("settings"):
+        settings = get_settings()
+        kw_list = settings.get("keywords", [])
+        if keyword in kw_list:
+            return False
+        kw_list.append(keyword)
+        settings["keywords"] = kw_list
+        _write_json("settings.json", settings)
+        return True
+
+
+def remove_keyword(keyword: str) -> bool:
+    """删除关键词，返回是否成功"""
+    with _get_lock("settings"):
+        settings = get_settings()
+        kw_list = settings.get("keywords", [])
+        if keyword not in kw_list:
+            return False
+        kw_list.remove(keyword)
+        settings["keywords"] = kw_list
+        _write_json("settings.json", settings)
+        return True