__init__.py

import toml
import os
from astropy.table import Table
import socket
from csst_dfs_client import plan, level0, level1, level2, dag, catalog

CONFIG = toml.load(os.path.join(os.path.dirname(__file__), "config.toml"))


def check_port(ip, port, timeout=3):
    """
    # # 示例：检查 192.168.1.1 的 80 端口是否开放
    # print(check_port("192.168.1.1", 80))  # True/False
    """
    try:
        # 创建 Socket 连接（TCP）
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)  # 设置超时时间（秒）
        # 尝试连接
        result = sock.connect_ex((ip, port))
        # 返回状态
        if result == 0:  # 0 表示成功
            return True
        else:
            return False
    except Exception as e:
        print(f"Error: {e}")
        return False
    finally:
        sock.close()  # 确保关闭连接


class DFS:

    plan = plan
    level0 = level0
    level1 = level1
    level2 = level2
    dag = dag
    catalog = catalog

    def __init__(self, location=None):
        # try each location
        print("Test all locations...", end="")
        status_list = []
        for loc in CONFIG.keys():
            dfs_ip = CONFIG[loc]["CSST_DFS_GATEWAY"].split(":")[0]
            dfs_port = int(CONFIG[loc]["CSST_DFS_GATEWAY"].split(":")[1])
            # redis_ip = CONFIG[loc]["redis"]["host"]
            # redis_port = CONFIG[loc]["redis"]["port"]
            this_dfs_status = check_port(dfs_ip, dfs_port, timeout=1)
            # this_redis_status = check_port(redis_ip, redis_port, timeout=1)
            this_status = dict(
                location=loc,
                status=this_dfs_status,  # and this_redis_status,
                dfs=this_dfs_status,
                # redis=this_redis_status,
            )
            status_list.append(this_status)
            # print(this_status)
        print("Done!\n")
        status_table = Table(status_list)
        print(status_table)
        print("\n")
        if status_table["status"].sum() == 0:
            print("No DFS location is available")
        elif status_table["status"].sum() > 1:
            print("Multiple DFS locations are available, please specify one")
        elif location is None:
            # set DFS automatically
            if status_table["status"].sum() == 1:
                print("One DFS locations are available, good")
                location = status_table["location"][status_table["status"]][0]
            elif status_table["status"].sum() == 0:
                print("No DFS location is available, using csu")
                location = "csu"
            else:
                raise ValueError("Multiple DFS locations are available")

            self.location = location
            self.config = CONFIG[self.location]

            for k, v in CONFIG[self.location].items():
                os.environ.setdefault(k, str(v))
            # print("Setting redis config:")
            # self.redis = Redis(location=self.location)


dfs = DFS(location=None)