核心功能:
- ssh_tunnel.py: SSH隧道管理器(连接香港ECS)
- tushare_source.py: A股数据获取(指数、ETF、期货)
- yfinance_source.py: 境外数据获取(港股、美股)
- hybrid_source.py: 混合数据源(整合所有)
使用方式:
from datasource import HybridDataSource
source = HybridDataSource.from_yaml('config/strategies/rotation.yaml')
result = source.fetch_all()
更新 RotationStrategy 使用新数据源模块
116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
"""
|
||
SSH隧道管理器
|
||
|
||
通过SSH隧道建立本地SOCKS5代理,用于访问境外数据源
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import time
|
||
import subprocess
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
|
||
class SSHTunnelManager:
|
||
"""SSH隧道管理器"""
|
||
|
||
def __init__(self, config: dict):
|
||
"""
|
||
初始化SSH隧道
|
||
|
||
Args:
|
||
config: SSH配置字典
|
||
- host: SSH服务器地址
|
||
- port: SSH端口(默认22)
|
||
- username: SSH用户名
|
||
- key_path: SSH私钥路径(相对或绝对)
|
||
- local_port: 本地SOCKS5端口(默认1080)
|
||
"""
|
||
self.enabled = config.get("enabled", False)
|
||
self.host = config.get("host", "")
|
||
self.port = config.get("port", 22)
|
||
self.username = config.get("username", "root")
|
||
self.local_port = config.get("local_port", 1080)
|
||
self._process: Optional[subprocess.Popen] = None
|
||
|
||
# 处理 key_path:相对路径转换为绝对路径
|
||
key_path = config.get("key_path", "")
|
||
if key_path and not os.path.isabs(key_path):
|
||
# 相对于项目根目录
|
||
project_root = Path(__file__).parent.parent
|
||
key_path = str(project_root / key_path)
|
||
self.key_path = key_path
|
||
|
||
def start(self) -> bool:
|
||
"""启动SSH隧道"""
|
||
if not self.enabled:
|
||
return True
|
||
|
||
if not all([self.host, self.username, self.key_path]):
|
||
print("SSH配置不完整,跳过隧道建立")
|
||
return False
|
||
|
||
print(f"建立SSH隧道: {self.host}:{self.port} -> 本地SOCKS5端口 {self.local_port}")
|
||
|
||
cmd = [
|
||
"ssh", "-N", "-D", f"127.0.0.1:{self.local_port}",
|
||
"-o", "StrictHostKeyChecking=no",
|
||
"-o", "UserKnownHostsFile=/dev/null",
|
||
"-i", self.key_path,
|
||
"-p", str(self.port),
|
||
f"{self.username}@{self.host}"
|
||
]
|
||
|
||
try:
|
||
self._process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||
time.sleep(2)
|
||
|
||
if self._process.poll() is not None:
|
||
stdout, stderr = self._process.communicate()
|
||
print("✗ SSH隧道建立失败")
|
||
if stderr:
|
||
print(f"错误: {stderr.decode()}")
|
||
return False
|
||
|
||
# 设置代理环境变量
|
||
# 使用 socks5h:// 让代理服务器远程解析DNS,避免IPv6问题
|
||
proxy_url = f"socks5h://127.0.0.1:{self.local_port}"
|
||
os.environ["HTTP_PROXY"] = proxy_url
|
||
os.environ["HTTPS_PROXY"] = proxy_url
|
||
os.environ["ALL_PROXY"] = proxy_url
|
||
|
||
print(f"✓ SSH隧道已建立: {proxy_url}")
|
||
time.sleep(1)
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"✗ SSH隧道异常: {e}")
|
||
return False
|
||
|
||
def stop(self):
|
||
"""停止SSH隧道"""
|
||
if self._process:
|
||
self._process.terminate()
|
||
self._process.wait()
|
||
for key in ["HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"]:
|
||
os.environ.pop(key, None)
|
||
print("SSH隧道已关闭")
|
||
|
||
def __enter__(self):
|
||
self.start()
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
self.stop()
|
||
|
||
|
||
def create_ssh_tunnel_from_yaml(config_path: str) -> SSHTunnelManager:
|
||
"""从YAML配置创建SSH隧道"""
|
||
import yaml
|
||
|
||
with open(config_path, 'r', encoding='utf-8') as f:
|
||
config = yaml.safe_load(f)
|
||
|
||
ssh_config = config.get('ssh_tunnel', {})
|
||
return SSHTunnelManager(ssh_config) |