mirror of
https://github.com/kodxana/madiator-docker-runpod.git
synced 2024-12-12 17:26:33 +01:00
70 lines
2.3 KiB
Python
70 lines
2.3 KiB
Python
|
import os
|
||
|
|
import subprocess
|
||
|
|
import json
|
||
|
|
|
||
|
|
SSH_CONFIG_FILE = '/etc/ssh/sshd_config'
|
||
|
|
SSH_PASSWORD_FILE = '/workspace/.ssh_password'
|
||
|
|
|
||
|
|
def save_ssh_password(password):
|
||
|
|
with open(SSH_PASSWORD_FILE, 'w') as f:
|
||
|
|
json.dump({'password': password}, f)
|
||
|
|
|
||
|
|
def get_ssh_password():
|
||
|
|
if os.path.exists(SSH_PASSWORD_FILE):
|
||
|
|
with open(SSH_PASSWORD_FILE, 'r') as f:
|
||
|
|
data = json.load(f)
|
||
|
|
return data.get('password')
|
||
|
|
return None
|
||
|
|
|
||
|
|
def check_ssh_config():
|
||
|
|
try:
|
||
|
|
with open('/etc/ssh/sshd_config', 'r') as f:
|
||
|
|
config = f.read()
|
||
|
|
|
||
|
|
root_login = 'PermitRootLogin yes' in config.split('\n')
|
||
|
|
password_auth = 'PasswordAuthentication yes' in config.split('\n')
|
||
|
|
|
||
|
|
print(f"Root login enabled: {root_login}")
|
||
|
|
print(f"Password authentication enabled: {password_auth}")
|
||
|
|
|
||
|
|
return root_login, password_auth
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error checking SSH config: {e}")
|
||
|
|
return False, False
|
||
|
|
|
||
|
|
def setup_ssh():
|
||
|
|
try:
|
||
|
|
print("Setting up SSH configuration...")
|
||
|
|
|
||
|
|
# Check for SSH host keys
|
||
|
|
print("Checking for SSH host keys...")
|
||
|
|
if not os.path.exists('/etc/ssh/ssh_host_ed25519_key'):
|
||
|
|
print("SSH host keys not found. Generating new host keys...")
|
||
|
|
subprocess.run(['ssh-keygen', '-t', 'ed25519', '-f', '/etc/ssh/ssh_host_ed25519_key', '-N', '""'], check=True)
|
||
|
|
subprocess.run(['ssh-keygen', '-t', 'rsa', '-b', '4096', '-f', '/etc/ssh/ssh_host_rsa_key', '-N', '""'], check=True)
|
||
|
|
print("SSH host keys generated successfully.")
|
||
|
|
else:
|
||
|
|
print("SSH host keys are already present.")
|
||
|
|
|
||
|
|
# Check if PUBLIC_KEY is set
|
||
|
|
public_key = os.environ.get('PUBLIC_KEY', '').strip()
|
||
|
|
|
||
|
|
if public_key:
|
||
|
|
# Ensure the .ssh directory exists
|
||
|
|
os.makedirs('/root/.ssh', exist_ok=True)
|
||
|
|
|
||
|
|
# Add the public key to authorized_keys
|
||
|
|
with open('/root/.ssh/authorized_keys', 'w') as f:
|
||
|
|
f.write(public_key + '\n')
|
||
|
|
|
||
|
|
# Set correct permissions
|
||
|
|
os.chmod('/root/.ssh', 0o700)
|
||
|
|
os.chmod('/root/.ssh/authorized_keys', 0o600)
|
||
|
|
|
||
|
|
print("SSH Configuration Updated.")
|
||
|
|
|
||
|
|
return True
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error setting up SSH: {str(e)}")
|
||
|
|
return False
|