hy-app-monitor/ops/open_port_via_tat.py
2026-04-22 00:15:58 +08:00

218 lines
7.7 KiB
Python
Executable File

#!/usr/bin/env python3
from __future__ import annotations
import base64
import json
import os
import time
from pathlib import Path
from typing import Callable, TypeVar
from tencentcloud.common import credential
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
from tencentcloud.common.profile.client_profile import ClientProfile
from tencentcloud.common.profile.http_profile import HttpProfile
from tencentcloud.cvm.v20170312 import cvm_client, models as cvm_models
from tencentcloud.tat.v20201028 import models as tat_models, tat_client
ROOT = Path(__file__).resolve().parents[1]
DEFAULT_INSTANCE_ID = os.environ.get("DEPLOY_INSTANCE_ID", "ins-ntmpcd3a")
DEFAULT_REGION = os.environ.get("DEPLOY_REGION", "me-saudi-arabia")
DEFAULT_PORT = int(os.environ.get("OPEN_PORT", "6666"))
DEFAULT_CIDR = os.environ.get("OPEN_CIDR", "0.0.0.0/0")
T = TypeVar("T")
def env_required(name: str) -> str:
value = os.environ.get(name, "").strip()
if not value:
raise SystemExit(f"missing env: {name}")
return value
def build_cvm(secret_id: str, secret_key: str, region: str) -> cvm_client.CvmClient:
cred = credential.Credential(secret_id, secret_key)
return cvm_client.CvmClient(
cred,
region,
ClientProfile(httpProfile=HttpProfile(endpoint="cvm.tencentcloudapi.com")),
)
def build_tat(secret_id: str, secret_key: str, region: str) -> tat_client.TatClient:
cred = credential.Credential(secret_id, secret_key)
return tat_client.TatClient(
cred,
region,
ClientProfile(httpProfile=HttpProfile(endpoint="tat.tencentcloudapi.com")),
)
def invoke(action: Callable[[], T], label: str, retries: int = 4, delay_seconds: int = 2) -> T:
last_error: Exception | None = None
for attempt in range(1, retries + 1):
try:
return action()
except TencentCloudSDKException as exc:
last_error = exc
if attempt == retries:
break
time.sleep(delay_seconds * attempt)
if last_error is not None:
raise last_error
raise RuntimeError(f"{label} failed without error")
def resolve_security_group_id(client: cvm_client.CvmClient, instance_id: str) -> str:
req = cvm_models.DescribeInstancesRequest()
req.from_json_string(json.dumps({"InstanceIds": [instance_id]}))
resp = json.loads(invoke(lambda: client.DescribeInstances(req), "DescribeInstances").to_json_string())
instances = resp.get("InstanceSet") or []
if not instances:
raise SystemExit(f"instance not found: {instance_id}")
security_group_ids = instances[0].get("SecurityGroupIds") or []
if not security_group_ids:
raise SystemExit(f"no security group attached: {instance_id}")
return str(security_group_ids[0])
def remote_script(secret_id: str, secret_key: str, region: str, security_group_id: str, port: int, cidr: str) -> str:
payload = {
"secret_id": secret_id,
"secret_key": secret_key,
"region": region,
"security_group_id": security_group_id,
"port": port,
"cidr": cidr,
}
blob = base64.b64encode(json.dumps(payload).encode("utf-8")).decode("ascii")
return f"""set -euo pipefail
if command -v firewall-cmd >/dev/null 2>&1; then
firewall-cmd --add-port={port}/tcp --permanent >/dev/null 2>&1 || true
firewall-cmd --reload >/dev/null 2>&1 || true
fi
if command -v ufw >/dev/null 2>&1; then
ufw allow {port}/tcp >/dev/null 2>&1 || true
fi
if command -v iptables >/dev/null 2>&1; then
iptables -C INPUT -p tcp --dport {port} -j ACCEPT >/dev/null 2>&1 || iptables -I INPUT -p tcp --dport {port} -j ACCEPT || true
fi
python3 -m pip install --user -q tencentcloud-sdk-python >/dev/null 2>&1 || true
python3 - <<'PY'
import base64
import json
from tencentcloud.common import credential
from tencentcloud.common.profile.client_profile import ClientProfile
from tencentcloud.common.profile.http_profile import HttpProfile
from tencentcloud.vpc.v20170312 import vpc_client, models
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
payload = json.loads(base64.b64decode({blob!r}).decode('utf-8'))
cred = credential.Credential(payload['secret_id'], payload['secret_key'])
client = vpc_client.VpcClient(
cred,
payload['region'],
ClientProfile(httpProfile=HttpProfile(endpoint='vpc.tencentcloudapi.com')),
)
req = models.CreateSecurityGroupPoliciesRequest()
req.from_json_string(json.dumps({{
'SecurityGroupId': payload['security_group_id'],
'SecurityGroupPolicySet': {{
'Ingress': [{{
'Protocol': 'TCP',
'Port': str(payload['port']),
'CidrBlock': payload['cidr'],
'Action': 'ACCEPT',
'PolicyDescription': 'hy-app-monitor 6666',
}}]
}}
}}))
try:
client.CreateSecurityGroupPolicies(req)
print('security_group_rule=created')
except TencentCloudSDKException as exc:
if 'Duplicate' in str(exc) or 'already exists' in str(exc):
print('security_group_rule=exists')
else:
raise
PY
echo "port_opened={port}"
"""
def run_tat_script(client: tat_client.TatClient, instance_id: str, script: str, timeout_seconds: int = 300) -> str:
req = tat_models.RunCommandRequest()
req.from_json_string(
json.dumps(
{
"CommandName": "hy-app-monitor-open-port",
"CommandType": "SHELL",
"Content": base64.b64encode(script.encode("utf-8")).decode("ascii"),
"InstanceIds": [instance_id],
"Timeout": timeout_seconds,
"WorkingDirectory": "/root",
}
)
)
resp = json.loads(invoke(lambda: client.RunCommand(req), "RunCommand").to_json_string())
invocation_id = resp["InvocationId"]
deadline = time.time() + timeout_seconds + 120
while time.time() < deadline:
time.sleep(3)
query = tat_models.DescribeInvocationTasksRequest()
query.from_json_string(
json.dumps(
{
"HideOutput": False,
"Limit": 20,
"Filters": [{"Name": "invocation-id", "Values": [invocation_id]}],
}
)
)
data = json.loads(invoke(lambda: client.DescribeInvocationTasks(query), "DescribeInvocationTasks").to_json_string())
tasks = data.get("InvocationTaskSet") or []
if not tasks:
continue
task = tasks[0]
status = task["TaskStatus"]
if status in {"SUCCESS", "FAILED", "TIMEOUT", "CANCELLED", "START_FAILED"}:
output = base64.b64decode(task.get("TaskResult", {}).get("Output", "")).decode("utf-8", errors="replace")
if status != "SUCCESS":
raise SystemExit(output or f"TAT task failed: {status}")
return output
raise SystemExit("TAT task timed out")
def main() -> None:
secret_id = env_required("TENCENT_SECRET_ID")
secret_key = env_required("TENCENT_SECRET_KEY")
instance_id = os.environ.get("DEPLOY_INSTANCE_ID", DEFAULT_INSTANCE_ID)
region = os.environ.get("DEPLOY_REGION", DEFAULT_REGION)
port = int(os.environ.get("OPEN_PORT", str(DEFAULT_PORT)))
cidr = os.environ.get("OPEN_CIDR", DEFAULT_CIDR)
cvm = build_cvm(secret_id, secret_key, region)
tat = build_tat(secret_id, secret_key, region)
security_group_id = os.environ.get("SECURITY_GROUP_ID", "").strip() or resolve_security_group_id(cvm, instance_id)
output = run_tat_script(tat, instance_id, remote_script(secret_id, secret_key, region, security_group_id, port, cidr))
print(f"instance_id={instance_id}")
print(f"security_group_id={security_group_id}")
print(output.strip())
if __name__ == "__main__":
main()