#!/usr/bin/env python
#
# Copyright 20011 Canonical Ltd.
#
# Checkbox is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Checkbox is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Checkbox.  If not, see <http://www.gnu.org/licenses/>.

'''
This script allows automated testing of Wake-on-Lan support.
 It asks a suitably configured server for a "wakeup call" in a set
 number of seconds (say, 120). It then uses rtcwake to send itself
 to sleep with a scheduled wakeup in a number of seconds, higher 
 than the previous one (say, 300).
 The script then records time elapsed before waking up and uses that
 to determine whether WoL is working.

 Assumptions:
 - rtcwake works correctly and the system auto-wakes when rtcwakealarm 
   goes off.
 - There's a URL that can be called to ask the server to schedule the wakeup
 call. It should specify the number of seconds to sleep as a parameter. If
 there's a problem scheduling the WOL call, the server should return an HTTP
 error code (4xx/5xx); or 2xx/3xx (success) if the call was scheduled
 successfully. The called CGI is responsible for obtaining our MAC address
 and scheduling the sending of WOL packets at the right time.
 - I will be called as root, so I don't need to call sudo myself.
'''

from optparse import OptionParser
import logging
import os
import sys
import time
import urllib2

RTCWAKE_COMMAND= "/usr/sbin/rtcwake"

def call_wol_url(url):
    try:
        return urllib2.urlopen(url).code
    except urllib2.URLError as error:
        if hasattr(error,'code'):
            logging.error("HTTP error calling Wake on LAN URL: %s" % 
                    error.code)
        else:
            logging.error("Invalid Wake on LAN URL.")
        return False

def sleep_for(seconds=120):
    rtcwake_command = "%s -s %d -m mem" % (RTCWAKE_COMMAND, seconds)
    try:
        rtcwake_exit_code = os.system(rtcwake_command)
        if rtcwake_exit_code == 0:
            return True
        else:
            logging.error("rtcwake command returned an error")
    except:
        logging.error("Unable to execute the rtcwake command")
    
    return False

def main():
    usage = "Usage: %prog [OPTIONS] WAKEUP_URL"
    parser = OptionParser(usage)
    parser.add_option('-w','--wol',
            action='store',
            type='int',
            metavar='NUM',
            dest='wol_seconds',
            default=120,
            help="Time after which I expect to receive the Wake on LAN \
                  packet and wake up. If system wakes up this many \
                  seconds after going to sleep via rtcwake, it is \
                  assumed that WOL worked and success is returned. \
                  Must be smaller than --rtcwake time. \
                  Default is %default seconds")
    parser.add_option('-r','--rtcwake',
            action='store',
            type='int',
            metavar='NUM',
            dest='rtcwake_seconds',
            default=300,
            help="Time after which system should wake up from the \
                  internal rtc-scheduled wakeup clock. If system \
                  wakes up after this time, test is marked failed. \
                  Must be bigger than --wol time. \
                  Default is %default seconds")
    parser.add_option('-d','--debug',
            action='store_true',
            default=False,
            help='Choose this to add verbose output for debug \
                  purposes')

    (options, args) = parser.parse_args()
    if not args:
        parser.error("Must specify the URL to schedule Wake on LAN packet")
        return 2

    if options.wol_seconds >= options.rtcwake_seconds:
        parser.error("rtcwake and wake on lan times make no sense")
        parser.print_help()
        return 2

    wake_url = args[0]

    if options.debug:
       logger = logging.getLogger()
       logger.setLevel(logging.DEBUG)

    if not call_wol_url(wake_url):
        return 2
    
    logging.info("Entering sleep for %d seconds" % options.rtcwake_seconds)
    fell_asleep_at = time.time()
    
    #invoke rtcwake here, if fail just exit the whole operation.
    if not sleep_for(seconds=options.rtcwake_seconds):
        return 2
    
    #I woke up, what time is it?
    awoke_at = time.time()

    #How long was I asleep?
    sleep_time = awoke_at - fell_asleep_at

    logging.info("Slept for %d seconds" % sleep_time)
    
    if sleep_time <= options.wol_seconds:
        logging.info( "Test FAILED (sleep duration too short)")
        return 1
    if sleep_time > options.wol_seconds and \
       sleep_time < options.rtcwake_seconds:
        logging.info("Test PASSED (sleep duration just right)")
        return 0
    if sleep_time >= options.rtcwake_seconds:
        logging.info("Test FAILED (sleep duration too long)")
        return 1

if __name__ == "__main__":
    sys.exit(main())
