Telegram 憑證監控機器人實作 EP5 - 代碼優化

{'subject': ((('countryName', 'US'),), (('stateOrProvinceName', 'California'),), (('localityName', 'San Francisco'),), (('organizationName', 'Wikimedia Foundation, Inc.'),), (('commonName', '*.wikipedia.org'),)), 'issuer': ((('countryName', 'US'),), (('organizationName', 'DigiCert Inc'),), (('commonName', 'DigiCert TLS Hybrid ECC SHA384 2020 CA1'),)), 'version': 3, 'serialNumber': '07419E39583A4C76CF1EA14347FA5F3A', 'notBefore': 'Oct 18 00:00:00 2023 GMT', 'notAfter': 'Oct 16 23:59:59 2024 GMT', 'subjectAltName': (('DNS', '*.wikipedia.org'), ('DNS', 'wikimedia.org'), ('DNS', 'mediawiki.org'), ('DNS', 'wikibooks.org'), ('DNS', 'wikidata.org'), ('DNS', 'wikinews.org'), ('DNS', 'wikiquote.org'), ('DNS', 'wikisource.org'), ('DNS',
....

def parse_ssl_cert_info(domain, cert):
    issuer = dict(x[0] for x in cert["issuer"])
    subject = dict(x[0] for x in cert["subject"])
    issued_to = subject.get("commonName", subject.get("organizationName", ""))
    issued_by = issuer.get("commonName", issuer.get("organizationName", ""))
    valid_from = datetime.strptime(cert["notBefore"], "%b %d %H:%M:%S %Y %Z").strftime(
        "%Y-%m-%d"
    )
    valid_until = datetime.strptime(cert["notAfter"], "%b %d %H:%M:%S %Y %Z").strftime(
        "%Y-%m-%d"
    )

    cert_info = (
        f"Domain: {domain}\n"
        f"Issued To: {issued_to}\n"
        f"Issued By: {issued_by}\n"
        f"Valid From: {valid_from}\n"
        f"Valid Until: {valid_until}"
    )
    return cert_info

@bot.message_handler(commands=["cert_info"])
def send_cert_info(message):
    _, get_domain = message.text.split(maxsplit=1)
    cert = get_ssl_cert_info(get_domain, check_only=False)
    cert_info = parse_ssl_cert_info(get_domain, cert)
    bot.send_message(message.chat.id, cert_info)

@bot.message_handler(commands=["cert_info"])
def send_cert_info(message):
    try:
        _, get_domain = message.text.split(maxsplit=1)
    except ValueError:
        # 處理使用者沒有提供域名的情況
        bot.send_message(message.chat.id, "請提供一個域名。例如：/cert_info example.com")
        return
    cert = get_ssl_cert_info(get_domain, check_only=False)
    if cert is None or cert is False:
        # 如果無法獲取證書，直接回傳錯誤訊息
        error_message = f"域名錯誤: 無法獲取 {get_domain} 的 SSL 證書資訊。"
        bot.send_message(message.chat.id, error_message)
    else:
        cert_info = parse_ssl_cert_info(get_domain, cert)
        bot.send_message(message.chat.id, cert_info)

@bot.message_handler(commands=["help"])
def send_help_message(message):
    help_text = """
以下是可用的命令列表及其用途：

/cert_info <domain> - 獲取指定域名的 SSL 證書資訊。
/get_all - 從 MongoDB 獲取所有域名環境的資訊。
/get <env> <domain> - 獲取指定環境下域名的資訊。
/edit <env> <old_domain> <new_domain> - 更新指定環境下域名的資訊。
/add <env> <domain> - 向 MongoDB 添加一個新的域名及其環境。
/del <env> <domain> - 從 MongoDB 刪除指定環境下的域名。
/bulk_add <env> <domain1> <domain2> ... - 向 MongoDB 批量添加多個域名及其環境。

請根據需要使用上述命令。
"""
    bot.send_message(message.chat.id, help_text)

mongodb_uri: "mongodb://rootuser:rootpass@localhost:27017/mydatabase?authSource=admin"
telegram_bot_token: "<YOUR-BOT-TOKEN>"
platform: "gcp"
telegram_group_id: "<YOUR-CHAT-ID>"

# config_loader.py
import yaml

def load_config():
    with open("config.yaml", "r") as file:
        return yaml.safe_load(file)

def get_mongodb_uri():
    config = load_config()
    return config.get("mongodb_uri", None)

def get_telegram_config():
    config = load_config()
    return {
        "telegram_bot_token": config.get("telegram_bot_token", None),
        "platform": config.get("platform", None),
        "telegram_group_id": config.get("telegram_group_id", None)
    }

services
  telegram_bot:
    build: .
    container_name: telegram_bot_container
    depends_on:
      - mongodb
    environment:
      TELEGRAM_BOT_TOKEN: <YOUR-BOT-TOKEN>
      TELEGRAM_GROUP_ID: <YOUR-CHAT-ID>
      MONGODB_URI: mongodb://rootuser:rootpass@mongodb:27017/mydatabase?authSource=admin
    command: python main.py
    restart: always

import os

def get_mongodb_uri():
    return os.getenv('MONGODB_URI', 'mongodb://localhost:27017/mydatabase')

def get_telegram_config():
    return {
        "telegram_bot_token": os.getenv('TELEGRAM_BOT_TOKEN', ''),
        "platform": os.getenv('PLATFORM', 'local'),
        "telegram_group_id": os.getenv('TELEGRAM_GROUP_ID', '')
    }

from config_loader import get_telegram_config

telegram_config = get_telegram_config()
telegram_bot_token = telegram_config["telegram_bot_token"]
platform = telegram_config["platform"]
telegram_group_id = telegram_config["telegram_group_id"]