app.py

import os
import streamlit as st
from dateutil import parser
import googleapiclient.discovery
from datetime import datetime, timedelta
from google.oauth2 import service_account


def auth():
    with open("/tmp/token.json", "w") as token_f:
        google_key = os.getenv("GOOGLE_KEY")
        token_f.write(google_key)
        
    credentials = service_account.Credentials.from_service_account_file('/tmp/token.json')
    
    return credentials

def get_status(credentials, instance_name):
    status = False
    client = googleapiclient.discovery.build('compute', 'v1', credentials=credentials)
    response = client.instances().list(project='nus-cisco-corp-lab-wp1', zone='asia-southeast1-c').execute()
    for item in response['items']:
        if item["name"] == instance_name:
            try:
                ip_address = item['networkInterfaces'][0]['accessConfigs'][0]['natIP']
            except Exception as e:
                ip_address = None
            status = True if item["status"] == "RUNNING" else False
    return status, ip_address

def get_time_remaining(last_start_timestamp, run_duration_seconds):
    if run_duration_seconds == -1:
        return -1
    run_duration_seconds = int(run_duration_seconds)
    last_start_dt = parser.isoparse(last_start_timestamp)
    end_run_dt = last_start_dt + timedelta(seconds=run_duration_seconds)
    now_dt = datetime.now(tz=last_start_dt.tzinfo)
    delta = int((end_run_dt - now_dt).total_seconds() / 60)
    return max(delta, 0)

def get_server(credentials, instance_name):
    service = googleapiclient.discovery.build('compute', 'v1', credentials=credentials)
    request = service.instances().get(project='nus-cisco-corp-lab-wp1', zone='asia-southeast1-c', instance=instance_name)
    response = request.execute()
    return response

def schedule_server(credentials, instance_name, running_time):
    service = googleapiclient.discovery.build('compute', 'v1', credentials=credentials)
    request = service.instances().setScheduling(project='nus-cisco-corp-lab-wp1', zone='asia-southeast1-c', instance=instance_name, body={"automaticRestart": False,"maxRunDuration": {"seconds": running_time*3600}})    
    response = request.execute()
    return response

def activate_server(credentials, instance_name):
    service = googleapiclient.discovery.build('compute', 'v1', credentials=credentials)
    request = service.instances().start(project='nus-cisco-corp-lab-wp1', zone='asia-southeast1-c', instance=instance_name)
    response = request.execute()
    return response

def deactivate_server(credentials, instance_name):
    service = googleapiclient.discovery.build('compute', 'v1', credentials=credentials)
    request = service.instances().stop(project='nus-cisco-corp-lab-wp1',zone='asia-southeast1-c',instance=instance_name, discardLocalSsd=False)
    response = request.execute()
    return response

st.title("Lucky Reactor Panel")

instance_name = "lucky-reactor"
credentials = auth()

# Get instance info
instance_info = get_server(credentials, instance_name)
status = instance_info['status']
ip_address = instance_info['networkInterfaces'][0]['accessConfigs'][0]['natIP'] if status == 'RUNNING' else None
max_run_duration = instance_info['scheduling']['maxRunDuration']['seconds'] if 'maxRunDuration' in instance_info['scheduling'] else -1
last_start_timestamp = instance_info['lastStartTimestamp']
max_remaining_time = get_time_remaining(last_start_timestamp, max_run_duration)

# UI
st.image("cover.jpg", caption="Lucky Reactor is currently " + ("running 🚀" if status == 'RUNNING' else "sleeping 😴") + ".")

if status == 'TERMINATED':            
    with st.form("activate_form"):
        token = st.text_input("**Token** (The reactor costs USD $10 per hour to operate)")
        running_time = st.select_slider("**Running Hours** (How many hours do you need the reactor)", options=[1,2,3,4,5,6])
        submitted = st.form_submit_button("Ignite 🚀")
        
        if submitted and token == os.getenv("EASY_TOKEN"):
            st.write("Lucky Reactor has been ignited. Please wait a few minutes (3~5) to preceed...")     
            schedule_response = schedule_server(credentials, instance_name, running_time)
            activate_response = activate_server(credentials, instance_name)
elif status == 'RUNNING':
    st.write(f"You can access Lucky Reactor via http://{ip_address}:7860")
    st.write(f'The reactor can still run **{max_remaining_time}** minutes. You can terminate it in advance.')
    
    with st.form("deactivate_form"):
        token = st.text_input("Token")
        terminated_btn = st.form_submit_button("Terminate 🕊️")
        
        if terminated_btn and token == os.getenv("EASY_TOKEN"):
            st.write("Lucky Reactor has been terminated.")            
            deactivate_response = deactivate_server(credentials, instance_name)
else:
    st.write(f"Server is [{status}]. Please wait...")