#!/usr/bin/python3
import logging
import argparse
import random
import time
import signal
import sys
import toml

import pwnagotchi
from pwnagotchi import utils
from pwnagotchi.plugins import cmd as plugins_cmd
from pwnagotchi import log
from pwnagotchi import restart
from pwnagotchi import fs
from pwnagotchi.utils import DottedTomlEncoder
from pwnagotchi._version import __version__ as _pwnagotchi_version

def _print_banner():
    logging.info(" ═══════════════ (⌐■_■) pwnagotchi64 v%s started ═══════════════" % _pwnagotchi_version)


def do_clear(display):
    logging.info("clearing the display ...")
    display.clear()
    sys.exit(0)


def do_show_message(display, text):
    # e-ink holds whatever was last rendered with zero power once this
    # process exits -- lets an external caller (e.g. pwnagotchi-syswatchdog,
    # which isn't part of the Python process and can't reach into a running
    # Agent's view) put a message on screen and have it stick around after
    # the real pwnagotchi service is stopped, same idea as agent.py's
    # set_rebooting()/on_shutdown() but usable standalone
    logging.info("showing message: %s" % text)
    import pwnagotchi.ui.faces as faces
    display.set('face', faces.SLEEP)
    display.set('status', text)
    display.update(force=True)
    time.sleep(3)  # give the e-ink time to finish its refresh before exiting
    sys.exit(0)


def do_manual_mode(agent):
    logging.info("entering manual mode ...")

    agent.mode = 'manual'
    agent.last_session.parse(agent.view(), args.skip_session)
    if not args.skip_session:
        logging.info(
            "the last session lasted %s (%d completed epochs, trained for %d), average reward:%s (min:%s max:%s)" % (
                agent.last_session.duration_human,
                agent.last_session.epochs,
                agent.last_session.train_epochs,
                agent.last_session.avg_reward,
                agent.last_session.min_reward,
                agent.last_session.max_reward))

    while True:
        display.on_manual_mode(agent.last_session)
        time.sleep(5)
        if grid.is_connected():
            plugins.on('internet_available', agent)


def do_auto_mode(agent):
    logging.info("entering auto mode ...")

    agent.start()

    while True:
        try:
            agent.recon()
            channels = agent.get_access_points_by_channel()
            for ch, aps in channels[:1]:
                agent.set_channel(ch)

                if not agent.is_stale() and agent.any_activity():
                    logging.info("%d access points on channel %d" % (len(aps), ch))

                for ap in aps:
                    agent.associate(ap)
                    for sta in ap['clients']:
                        agent.deauth(ap, sta)

            agent.flush_pending_wait()

            if agent.mode == 'auto' and not agent.is_training() and not agent.any_activity():
                personality_channels = agent._config['personality']['channels']
                wait_channel = random.choice(personality_channels) if personality_channels \
                    else random.choice([1, 6, 11])
                hop_recon_time = agent._config['personality']['hop_recon_time']
                logging.info("nothing to attack this epoch -- holding on channel %d for %ds anyway ...",
                              wait_channel, hop_recon_time)
                agent.set_channel(wait_channel)
                agent.wait_for(hop_recon_time)

            agent.next_epoch()

            if grid.is_connected():
                plugins.on('internet_available', agent)

        except Exception as e:
            if str(e).find("wifi.interface not set") > 0:
                logging.exception("main loop exception due to unavailable wifi device, likely programmatically disabled (%s)", e)
                logging.info("sleeping 60 seconds then advancing to next epoch to allow for cleanup code to trigger")
                time.sleep(60)
                agent.next_epoch()
            else:
                logging.exception("main loop exception (%s)", e)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser = plugins_cmd.add_parsers(parser)

    parser.add_argument('-C', '--config', action='store', dest='config', default='/etc/pwnagotchi/default.toml',
                        help='Main configuration file.')
    parser.add_argument('-U', '--user-config', action='store', dest='user_config', default='/etc/pwnagotchi/config.toml',
                        help='If this file exists, configuration will be merged and this will override default values.')

    parser.add_argument('--manual', dest="do_manual", action="store_true", default=False, help="Manual mode.")
    parser.add_argument('--skip-session', dest="skip_session", action="store_true", default=False,
                        help="Skip last session parsing in manual mode.")

    parser.add_argument('--clear', dest="do_clear", action="store_true", default=False,
                        help="Clear the ePaper display and exit.")

    parser.add_argument('--show-message', dest="show_message", action="store", default=None,
                        help="Show a sleep-face message on the display and exit "
                             "(sticks around after the process exits, since e-ink "
                             "needs no power to hold a frame).")

    parser.add_argument('--debug', dest="debug", action="store_true", default=False,
                        help="Enable debug logs.")

    parser.add_argument('--version', dest="version", action="store_true", default=False,
                        help="Print the version.")

    parser.add_argument('--print-config', dest="print_config", action="store_true", default=False,
                        help="Print the configuration.")

    parser.add_argument('--setup', dest="do_setup", action="store_true", default=False,
                        help="Run the interactive configuration setup wizard and exit.")

    args = parser.parse_args()

    if args.do_setup:
        from pwnagotchi.setup_wizard import run_wizard
        sys.exit(run_wizard(args))

    if plugins_cmd.used_plugin_cmd(args):
      config = utils.load_config(args)
      log.setup_logging(args, config)
      rc = plugins_cmd.handle_cmd(args, config)
      sys.exit(rc)

    if args.version:
        print(pwnagotchi.__version__)
        sys.exit(0)

    config = utils.load_config(args)

    if args.print_config:
        print(toml.dumps(config, encoder=DottedTomlEncoder()))
        sys.exit(0)

    from pwnagotchi.identity import KeyPair
    from pwnagotchi.agent import Agent
    from pwnagotchi.ui import fonts
    from pwnagotchi.ui.display import Display
    from pwnagotchi import grid
    from pwnagotchi import plugins

    pwnagotchi.config = config
    fs.setup_mounts(config)
    log.setup_logging(args, config)
    _print_banner()
    fonts.init(config)

    pwnagotchi.set_name(config['main']['name'])

    plugins.load(config)

    display = Display(config=config, state={'name': pwnagotchi.name()})

    if args.do_clear:
        do_clear(display)
        sys.exit(0)

    if args.show_message:
        do_show_message(display, args.show_message)
        sys.exit(0)

    agent = Agent(view=display, config=config, keypair=KeyPair(view=display))

    def usr1_handler(*unused):
        logging.info('Received USR1 singal. Restart process ...')
        restart("MANU" if args.do_manual else "AUTO")

    signal.signal(signal.SIGUSR1, usr1_handler)

    def sigterm_handler(*unused):
        try:
            agent._save_recovery_data()
        except Exception as e:
            logging.warning("couldn't save recovery data on SIGTERM: %s", e)
        sys.exit(0)

    signal.signal(signal.SIGTERM, sigterm_handler)

    if args.do_manual:
        do_manual_mode(agent)
    else:
        do_auto_mode(agent)
