main-py

main.py

import MySQLdb
import os
import sys
from sqlalchemy import create_engine, MetaData
from sqlalchemy.sql import text
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy import Column, String, Integer, Date, Table, ForeignKey, Float
from sqlalchemy.orm import relationship
from sqlalchemy.ext.declarative import declarative_base
from urllib import quote_plus as urlquote
from multiprocessing.dummy import Pool as ThreadPool 
import json
import urllib
import smtplib
from email.mime.text import MIMEText
import requests
import paramiko
import time
from config import *
import boto3
from models import Subscription, Server

# Set PID file, this prevents the script from running if already running
pid = str(os.getpid())
pidfile = "/tmp/main-py.pid"
if os.path.isfile(pidfile):
    print "%s already exists, exiting" % pidfile
    sys.exit()
file(pidfile, 'w').write(pid)

# Setup SQLAlchemy
engine = create_engine('mysql://ghostifidbuser:%s@localhost:3306/ghostifi' % urlquote(DB_PASSWORD), echo=False)
metadata = MetaData(bind=engine)
Session = scoped_session(sessionmaker(engine, autoflush=True))
session = Session()
Base = declarative_base()
Base.metadata.create_all(engine)

try:
    # Create lists which will be populated by SQL queries
    servers_to_create = []
    servers_to_destroy = []

    # Get all active subscriptions, joined to servers
    active_subscriptions = session.query(Subscription, Server).outerjoin(Server, Subscription.id == Server.wp_edd_sub_id).filter(Subscription.status=="Active").all()
    
    # Find active subscriptions which do not have existing servers (create these)
    for subscription in active_subscriptions:
        # If subscription exists for this server, skip, else append to the server create list
        if subscription[1]:
            pass
        else:
            servers_to_create.append(subscription[0])

    # Get all existing servers, joined to subscriptions
    active_servers = session.query(Server, Subscription).outerjoin(Subscription, Subscription.id == Server.wp_edd_sub_id).all()

    # Find existing servers which do not have active subscriptions (destroy these)
    for server in active_servers:
        # If subscription exists for this server, skip, else append to the server destroy list
        if server[1]:
            pass
        else:
            servers_to_destroy.append(server[0])

    # Get servers marked for delete
    delete_request_servers = session.query(Server, Subscription).outerjoin(Subscription, Subscription.id == Server.wp_edd_sub_id).filter(Server.delete_request == 1).all()

    for server in delete_request_servers:
        servers_to_destroy.append(server[0])

    # Get all servers which need to be rebuilt now (rebuild these)
    servers_to_rebuild_now = session.query(Server).filter(Server.rebuild_now_status==1).all()

    # Make the Pool of workers
    pool = ThreadPool(4)

    def thread_servers_to_create(self):
        # Setup SQLAlchemy
        session = Session()
        # If product is The VPS VPN Monthly or Yearly, set the VPS plan_id to 201 (25GB SSD/1GB RAM)
        if self.product_id == 5747 or self.product_id == 5745:
            vps_plan_id = '201'
            bandwidth_limit_this_month = 1000.00

        # Get username, email out of wp_users table. Using a raw query because this is a one-off
        sql = text('Select wp_users.user_login, wp_users.user_email from wp_edd_subscriptions left join wp_edd_customers on wp_edd_subscriptions.customer_id = wp_edd_customers.id left join wp_users on wp_edd_customers.user_id = wp_users.id where wp_edd_subscriptions.id = :id')
        query = session.execute(sql, {'id': self.id}).fetchone()
        username = query[0]
        email = query[1]
        # Create a new server object
        new_server = Server(self.customer_id, self.product_id, self.id, '0.0.0.0', '', email, '', '', 'Installing', bandwidth_limit_this_month, vps_plan_id, username, '')
        # Save to database
        session.add(new_server)
        session.commit()
        new_server.create()
        session.commit()
        # Close session
        session.close()

    def thread_servers_to_destroy(self):
        # Change status to Rebuilding to display this in the user dashboard
        self.status = 'Destroying'
        session.add(self)
        # Update database set status from Running to Rebuilding
        session.commit()
        # Run rebuild function
        self.destroy()
        # Commit changes - update status, rebuild_location, bandwidth_this_month etc
        session.delete(self)
        session.commit()

    def thread_servers_to_rebuild_now(self):
        # Change status to Rebuilding to display this in the user dashboard
        self.status = 'Rebuilding'
        self.rebuild_now_status = 0
        session.add(self)
        # Update database set status from Running to Rebuilding
        session.commit()
        # Run rebuild function
        self.rebuild('rebuild_now')
        # Commit changes - update status, rebuild_location, bandwidth_this_month etc
        session.add(self)
        session.commit()

    # Start a thread for each servers_to_create
    results = pool.map(thread_servers_to_create, servers_to_create)

    # Start a thread for each servers_to_destroy
    results = pool.map(thread_servers_to_destroy, servers_to_destroy)

    # Start a thread for each servers_to_rebuild_now
    results = pool.map(thread_servers_to_rebuild_now, servers_to_rebuild_now)

    # close the pool and wait for the work to finish
    pool.close()
    pool.join()
    session.close()

finally:
    os.unlink(pidfile)