#!/usr/bin/env python

import argparse
import telnetlib
import getpass

CMD_SYSTEM_MAINTENANCE = 24
CMD_DIAGNOSTIC = 4
CMD_RESET_XDSL = 1
CMD_REBOOT_SYSTEM = 21

class Router():
    def __init__(self, host, timeout=5):
        self.timeout = timeout
        self.tc = telnetlib.Telnet(host, timeout=timeout)

    def login(self, password):
        self._command('Password: ', password)

    def reboot(self):
        cmds = [CMD_SYSTEM_MAINTENANCE, CMD_DIAGNOSTIC, CMD_REBOOT_SYSTEM]
        [self._menu(c) for c in cmds]

    def reset(self):
        cmds = [CMD_SYSTEM_MAINTENANCE, CMD_DIAGNOSTIC, CMD_RESET_XDSL]
        [self._menu(c) for c in cmds]

    def _command(self, prompt, command):
        self.tc.read_until(prompt, self.timeout)
        self.tc.write(str(command) + '\n')

    def _menu(self, menu):
        self._command('Enter Menu Selection Number: ', menu)

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.tc.close()

def setup_parser():
# TODO write usage
# TODO add config reader
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', dest='config', default='~/.router',
                       help='config file destination')
    parser.add_argument('-p', '--password', dest='password',
                        help='you router admin password')
    parser.add_argument('-t', '--timeout', dest='timeout', default=3, type=int,
                        help='timeout in seconds')
    parser.add_argument('host', help='host name or IP address')
    parser.add_argument('command', choices=['reboot', 'reset'],
                        help='command to execute')
    return parser

def main():
    parser = setup_parser()
    args = parser.parse_args()

    password = args.password or getpass.getpass()

    with Router(args.host, args.timeout) as r:
        r.login(password)
        if args.command == 'reboot':
            r.reboot()
        elif args.command == 'reset':
            r.reset()

if __name__ == '__main__':
    main()