proxstar/proxstar/tasks.py
2018-09-05 22:11:07 -04:00

167 lines
5.7 KiB
Python

import os
import time
import requests
import paramiko
import psycopg2
from flask import Flask
from rq import get_current_job
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from proxstar.db import *
from proxstar.util import *
from proxstar.mail import *
from proxstar.starrs import *
from proxstar.vnc import send_stop_ssh_tunnel
from proxstar.vm import VM, create_vm, clone_vm
from proxstar.user import User, get_vms_for_rtp
from proxstar.proxmox import connect_proxmox, get_pools
app = Flask(__name__)
if os.path.exists(
os.path.join(
app.config.get('ROOT_DIR', os.getcwd()), "config.local.py")):
config = os.path.join(
app.config.get('ROOT_DIR', os.getcwd()), "config.local.py")
else:
config = os.path.join(app.config.get('ROOT_DIR', os.getcwd()), "config.py")
app.config.from_pyfile(config)
def connect_db():
engine = create_engine(app.config['SQLALCHEMY_DATABASE_URI'])
Base.metadata.bind = engine
DBSession = sessionmaker(bind=engine)
db = DBSession()
return db
def connect_starrs():
starrs = psycopg2.connect(
"dbname='{}' user='{}' host='{}' password='{}'".format(
app.config['STARRS_DB_NAME'], app.config['STARRS_DB_USER'],
app.config['STARRS_DB_HOST'], app.config['STARRS_DB_PASS']))
return starrs
def create_vm_task(user, name, cores, memory, disk, iso):
with app.app_context():
job = get_current_job()
proxmox = connect_proxmox()
db = connect_db()
starrs = connect_starrs()
job.meta['status'] = 'creating VM'
job.save_meta()
vmid, mac = create_vm(proxmox, user, name, cores, memory, disk, iso)
job.meta['status'] = 'registering in STARRS'
job.save_meta()
register_starrs(starrs, name, app.config['STARRS_USER'], mac,
get_next_ip(starrs, app.config['STARRS_IP_RANGE']))
job.meta['status'] = 'setting VM expiration'
job.save_meta()
delete_vm_expire(db, vmid)
get_vm_expire(db, vmid, app.config['VM_EXPIRE_MONTHS'])
job.meta['status'] = 'complete'
job.save_meta()
def delete_vm_task(vmid):
with app.app_context():
db = connect_db()
starrs = connect_starrs()
vm = VM(vmid)
vm.delete()
delete_starrs(starrs, vm.name)
delete_vm_expire(db, vmid)
def process_expiring_vms_task():
with app.app_context():
proxmox = connect_proxmox()
db = connect_db()
starrs = connect_starrs()
pools = get_pools(proxmox, db)
expired_vms = []
for pool in pools:
user = User(pool)
expiring_vms = []
vms = user.vms
for vm in vms:
vm = VM(vm['vmid'])
days = (vm.expire - datetime.date.today()).days
if days in [10, 7, 3, 1, 0, -1, -2, -3, -4, -5, -6]:
name = vm.name
expiring_vms.append([vm.id, vm.name, days])
if days <= 0:
expired_vms.append([vm.id, vm.name, days])
if days <= 0:
vm.stop()
elif days == -7:
print(
"Deleting {} ({}) as it has been a week since expiration.".
format(vm.name, vm.id))
send_stop_ssh_tunnel(vm.id)
delete_vm_task(vm.id)
if expiring_vms:
send_vm_expire_email(pool, expiring_vms)
if expired_vms:
send_rtp_vm_delete_email(expired_vms)
def generate_pool_cache_task():
with app.app_context():
proxmox = connect_proxmox()
db = connect_db()
pools = get_vms_for_rtp(proxmox, db)
store_pool_cache(db, pools)
def setup_template_task(template_id, name, user, ssh_key, cores, memory):
with app.app_context():
job = get_current_job()
proxmox = connect_proxmox()
starrs = connect_starrs()
db = connect_db()
print("[{}] Retrieving template info for template {}.".format(
name, template_id))
template = get_template(db, template_id)
print("[{}] Cloning template {}.".format(name, template_id))
job.meta['status'] = 'cloning template'
job.save_meta()
vmid, mac = clone_vm(proxmox, template_id, name, user)
print("[{}] Registering in STARRS.".format(name))
job.meta['status'] = 'registering in STARRS'
job.save_meta()
ip = get_next_ip(starrs, app.config['STARRS_IP_RANGE'])
register_starrs(starrs, name, app.config['STARRS_USER'], mac, ip)
get_vm_expire(db, vmid, app.config['VM_EXPIRE_MONTHS'])
print("[{}] Setting CPU and memory.".format(name))
job.meta['status'] = 'setting CPU and memory'
job.save_meta()
vm = VM(vmid)
vm.set_cpu(cores)
vm.set_mem(memory)
print("[{}] Applying cloud-init config.".format(name))
job.meta['status'] = 'applying cloud-init'
vm.set_ci_user(user)
vm.set_ci_ssh_key(ssh_key)
vm.set_ci_network()
print(
"[{}] Waiting for STARRS to propogate before starting VM.".format(
name))
job.meta['status'] = 'waiting for STARRS'
job.save_meta()
time.sleep(90)
print("[{}] Starting VM.".format(name))
job.meta['status'] = 'starting VM'
job.save_meta()
vm.start()
print("[{}] Template successfully provisioned.".format(name))
job.meta['status'] = 'completed'
job.save_meta()
def cleanup_vnc_task():
requests.post(
"https://{}/console/cleanup".format(app.config['SERVER_NAME']),
data={'token': app.config['VNC_CLEANUP_TOKEN']},
verify=False)