1
import uuid
2
import base64
3
import shutil
4
import subprocess
5
from pathlib import Path
6
from textwrap import dedent
7
from pyfiglet import Figlet
8

9
# --- 配置区 ---
10
OUTPUT_BASE_DIR = Path("./cloudinit_output")
11
HOSTNAMES = [
12
    "ubuntu-master",
13
    "ubuntu-worker01",
14
    "ubuntu-worker02"
15
]
16
# --- 配置区结束 ---
17

18
def generate_ssh_keypair(key_path: Path):
19
    key_path.parent.mkdir(parents=True, exist_ok=True)
20
    if key_path.exists():
21
        print(f"密钥文件已存在，跳过生成：{key_path}")
22
        return
23
    print(f"生成新的 SSH 密钥对: {key_path}")
24
    try:
25
        subprocess.run(
26
            ["ssh-keygen", "-t", "rsa", "-b", "4096", "-f", str(key_path), "-N", "", "-q"],
27
            check=True
28
        )
29
        print(f"密钥对已生成: {key_path} 和 {key_path.with_suffix('.pub')}")
30
    except (FileNotFoundError, subprocess.CalledProcessError) as e:
31
        print(f"[X] ssh-keygen 命令执行失败: {e}")
32
        print("[i] 请确保您的系统中安装了 OpenSSH 客户端，并且 'ssh-keygen' 在系统 PATH 中。")
33
        raise
34

35

36
def get_ssh_pubkey(pubkey_path: Path) -> str:
37
    if not pubkey_path.exists():
38
        raise FileNotFoundError(f"公钥文件不存在: {pubkey_path}")
39
    return pubkey_path.read_text(encoding="utf-8").strip()
40

41

42
def manual_input_pubkey() -> str:
43
    while True:
44
        print("\n请输入 SSH 公钥（单行格式，如 ssh-rsa AAAA...），输入完成后按回车：")
45
        key = input("公钥: ").strip()
46
        if key.startswith("ssh-rsa") or key.startswith("ssh-ed25519"):
47
            return key
48
        print("输入无效。公钥通常以 'ssh-rsa' 或 'ssh-ed25519' 开头。请重试。")
49

50

51
"""
52
        #cloud-config
53
        hostname: {hostname}
54
        fqdn: {hostname}.local
55
        manage_etc_hosts: true
56

57
        # 用户和认证设置
58
        disable_root: false
59
        ssh_pwauth: false
60

61
        users:
62
          - name: root
63
            # 关键修正：'ssh_authorized_keys' 必须与 'name' 对齐
64
            ssh_authorized_keys:
65
              - {ssh_key}
66

67
        # (可选) 设置 root 密码
68
        # chpasswd:
69
        #   list: |
70
        #     root:password123
71
        #   expire: false
72

73
        # 启动后执行的命令
74
        runcmd:
75
          # 1. 替换 APT 软件源为清华大学镜像 (推荐保留)
76
          - sed -i 's@https?://[^/]+/ubuntu/@https://mirrors.tuna.tsinghua.edu.cn/ubuntu/@g' /etc/apt/sources.list
77

78
          # 2. 增强 SSH 安全性，禁止 root 密码登录
79
          - sed -i 's/^#?PermitRootLogin.*/PermitRootLogin prohibit-password/' /etc/ssh/sshd_config
80

81
          # 3. 重启 SSH 服务使配置生效
82
          - systemctl restart sshd
83

84
          # 4. 更新系统
85
          - apt-get update
86
          - apt-get upgrade -y
87
    """
88
def generate_user_data(hostname: str, ssh_key: str) -> str:
89
    return dedent(f"""\
90
        #cloud-config
91
        hostname: {hostname}
92
        fqdn: {hostname}.local
93
        manage_etc_hosts: true
94
        disable_root: false
95
        ssh_pwauth: false
96
        users:
97
          - name: root
98
            ssh_authorized_keys:
99
              - {ssh_key}
100
        runcmd:
101
          - sed -i 's@https?://[^/]+/ubuntu/@https://mirrors.tuna.tsinghua.edu.cn/ubuntu/@g' /etc/apt/sources.list
102
          - sed -i 's/^#?PermitRootLogin.*/PermitRootLogin prohibit-password/' /etc/ssh/sshd_config
103
          - systemctl restart sshd
104
          - apt-get update
105
          - apt-get upgrade -y
106
    """)
107

108
def generate_meta_data(hostname: str) -> str:
109
    instance_id = f"{hostname}-{uuid.uuid4().hex[:8]}"
110
    return dedent(f"""\
111
        instance-id: {instance_id}
112
        local-hostname: {hostname}
113
    """)
114

115

116
def create_cloud_init_iso(hostname: str, source_dir: Path):
117
    command_name = "genisoimage"
118
    command_path = shutil.which(command_name)
119
    if not command_path:
120
        command_name = "mkisofs"
121
        command_path = shutil.which(command_name)
122

123
    if not command_path:
124
        print(f"[!] 警告: 未找到 'genisoimage' 或 'mkisofs' 命令，无法为 {hostname} 创建 seed.iso。")
125
        print("    - 在 Debian/Ubuntu 上, 请运行: sudo apt-get install -y cloud-image-utils")
126
        print("    - 在 CentOS/RHEL/Fedora 上, 请运行: sudo yum install -y genisoimage")
127
        print("    - 在 Windows 上, 您可以安装 cdrtools 或手动创建 ISO。")
128
        return
129

130
    iso_path = source_dir / "seed.iso"
131
    print(f"正在为 {hostname} 创建 cloud-init ISO: {iso_path} (使用 {command_name})")
132

133
    try:
134
        subprocess.run(
135
            [
136
                command_path,
137
                "-output", str(iso_path),
138
                "-volid", "cidata",
139
                "-joliet",
140
                "-rock",
141
                str(source_dir)
142
            ],
143
            check=True,
144
            capture_output=True
145
        )
146
        print(f"[√] 成功创建 {iso_path}")
147
    except subprocess.CalledProcessError as e:
148
        print(f"[X] 使用 {command_name} 创建 ISO 失败:")
149
        print(f"    错误码: {e.returncode}")
150
        print(f"    错误输出: {e.stderr.decode(errors='ignore')}")
151

152

153
def main():
154
    print("--- Cloud-Init 配置生成器 ---")
155
    ssh_pub_key = ""
156
    while True:
157
        choice = input(
158
            "请选择 SSH 公key获取方式:\n1) 手动输入公钥\n2) 自动生成/使用现有密钥对\n请输入 1 或 2: ").strip()
159
        if choice in ["1", "2"]: break
160
        print("无效选择，请重试。")
161
    try:
162
        if choice == "1":
163
            ssh_pub_key = manual_input_pubkey()
164
        else:
165
            priv_key_path = OUTPUT_BASE_DIR / "auth" / "id_rsa"
166
            generate_ssh_keypair(priv_key_path)
167
            ssh_pub_key = get_ssh_pubkey(priv_key_path.with_suffix(".pub"))
168
            print(f"\n[i] 将使用以下公钥（位于 {priv_key_path.with_suffix('.pub')}）:")
169
            print(ssh_pub_key)
170
    except Exception as e:
171
        print(f"\n[X] 获取SSH密钥时发生错误: {e}")
172
        return
173

174
    while True:
175
        out_choice = input(
176
            "\n请选择输出格式:\n1) 生成独立文件和 seed.iso (用于 NoCloud 数据源)\n2) 生成 Base64 键值对 (同时保存原始YAML文件)\n请输入 1 或 2: ").strip()
177
        if out_choice in ["1", "2"]: break
178
        print("无效选择，请重试。")
179

180
    print("\n--- 开始生成配置文件 ---")
181
    if out_choice == "1":
182
        for hostname in HOSTNAMES:
183
            host_dir = OUTPUT_BASE_DIR / hostname
184
            host_dir.mkdir(parents=True, exist_ok=True)
185
            user_data = generate_user_data(hostname, ssh_pub_key)
186
            meta_data = generate_meta_data(hostname)
187
            (host_dir / "user-data").write_text(user_data, encoding="utf-8")
188
            (host_dir / "meta-data").write_text(meta_data, encoding="utf-8")
189
            print(f"[√] 已为 {hostname} 生成配置文件于: {host_dir}/")
190
            create_cloud_init_iso(hostname, host_dir)
191
    else:  # out_choice == "2"
192
        all_configs_path = OUTPUT_BASE_DIR / "all_hosts_guestinfo.txt"
193
        all_configs_path.parent.mkdir(parents=True, exist_ok=True)
194
        with all_configs_path.open("w", encoding="utf-8") as f_all:
195
            for i, hostname in enumerate(HOSTNAMES):
196
                print(f"\n--- 正在处理: {hostname} ---")
197
                # 关键修复：使用 ssh_pub_key 而不是 ssh_key
198
                user_data_str = generate_user_data(hostname, ssh_pub_key)
199
                meta_data_str = generate_meta_data(hostname)
200
                host_dir = OUTPUT_BASE_DIR / hostname
201
                host_dir.mkdir(parents=True, exist_ok=True)
202
                (host_dir / "user-data.yaml").write_text(user_data_str, encoding="utf-8")
203
                (host_dir / "meta-data.yaml").write_text(meta_data_str, encoding="utf-8")
204
                print(f"[i] 原始YAML文件已保存至: {host_dir}/")
205
                user_data_b64 = base64.b64encode(user_data_str.encode("utf-8")).decode("utf-8")
206
                meta_data_b64 = base64.b64encode(meta_data_str.encode("utf-8")).decode("utf-8")
207
                output_content = dedent(f"""\
208
                    # Hostname: {hostname}
209
                    guestinfo.userdata="{user_data_b64}"
210
                    guestinfo.userdata.encoding="base64"
211
                    guestinfo.metadata="{meta_data_b64}"
212
                    guestinfo.metadata.encoding="base64"
213
                """)
214
                print(f"[√] 正在生成 Base64 配置...")
215
                if i > 0: f_all.write("\n")
216
                f_all.write(output_content)
217
        print(f"\n[√] 所有主机的 Base64 配置已合并保存到: {all_configs_path}")
218

219
    print("\n--- 所有任务完成 ---")
220

221

222
if __name__ == "__main__":
223
    try:
224
        f = Figlet(font='slant')
225
        print(f.renderText('cloud-init creat config by inuyume'))
226
    except NameError:
227
        print("--- cloud-init creat config by inuyume ---")
228

229
    main()
230

231
    print(
232
        "\n[Hint] : 如果想使用SSH密码，请在generate_user_data的dedent中修改prohibit-password为yes并取消chpasswd的注释并修改默认的password123")
233
    print(
234
        "\n[Hint] : 请注意保存生成的公钥和私钥和生成的配置文件，相同目录下运行脚本会覆盖旧的生成文件，密钥默认使用已有的密钥！")