/google/devshell/send_heartbeats.py

개요

/google/devshell/send_heartbeats.py

#!/usr/bin/python

# Script that polls the output of `who` to check if anyone is currently connected and sends a
# heartbeat if so.

import datetime,logging,os,requests,subprocess,time

logging.getLogger().setLevel(logging.INFO)

logging.info('Reading metadata...')

external_ip = os.environ['DEVSHELL_IP_ADDRESS']
logging.info('External IP address is %s' % external_ip)

def read_last_line(filename):
    with open(filename, 'r') as f:
        for line in f:
            pass
        return line

def send_heartbeat():
    url = os.environ['DEVSHELL_SERVER_BASE_URL'] + '/devshell/vmheartbeat?Devshell-Vm-Ip-Address=' + external_ip
    logging.info('Sending heartbeat to %s' % url)
    resp = requests.post(url, headers={"Devshell-Vm-Ip-Address": external_ip})
    logging.info('Got response code %d' % resp.status_code)

period = int(os.environ['DEVSHELL_HEARTBEAT_PERIOD_SECONDS'])
while True:
    time.sleep(period)
    should_send = False

    # who will print a line for each real terminal connected to the VM as well as
    # for each tmux session. The tmux sessions persist after the terminal is
    # closed, so filter them out.
    try:
        output = subprocess.check_output('who | grep -v \(tmux\(', shell=True)
        logging.info('who says:\n%s' % output)
        if len(output) > 0:
            should_send = True
    except Exception as e:
        logging.exception('who failed: %s' % e)

    # Also check auth.log to see when a user last logged in. This will catch
    # cases where users run one-off commands that don't result in sessions
    # visible to who.
    try:
        now = datetime.datetime.now()
        last_line = read_last_line('/var/log/auth.log')
        logging.info('auth.log says:\n%s', last_line)
        last_auth = datetime.datetime.strptime(filter(None, last_line.split(' '))[2], '%H:%M:%S')
        diff = (now-last_auth).seconds
        logging.info('last auth was %d seconds ago', diff)
        if diff < period:
            should_send = True
    except Exception as e:
        logging.exception('auth.log failed: %s' % e)

    try:
        if should_send:
            send_heartbeat()
    except Exception as e:
        logging.exception('heartbeat failed: %s' % e)