pwnagotchi/pwnagotchi/__init__.py

import logging
import time
import re
import argparse
import time
import signal
import sys
import toml
import requests
import os

from pwnagotchi._version import __version__

from pwnagotchi import utils
from pwnagotchi.google import cmd as google_cmd
from pwnagotchi.plugins import cmd as plugins_cmd
from pwnagotchi import log
from pwnagotchi import fs
from pwnagotchi.utils import DottedTomlEncoder, parse_version as version_to_tuple

_name = None
config = None
_cpu_stats = {}


def set_name(new_name):
    if new_name is None:
        return

    new_name = new_name.strip()
    if new_name == '':
        return

    if not re.match(r'^[a-zA-Z0-9\-]{2,25}$', new_name):
        logging.warning("name '%s' is invalid: min length is 2, max length 25, only a-zA-Z0-9- allowed", new_name)
        return

    current = name()
    if new_name != current:
        global _name

        logging.info("setting unit hostname '%s' -> '%s'", current, new_name)
        with open('/etc/hostname', 'wt') as fp:
            fp.write(new_name)

        with open('/etc/hosts', 'rt') as fp:
            prev = fp.read()
            logging.debug("old hosts:\n%s\n", prev)

        with open('/etc/hosts', 'wt') as fp:
            patched = prev.replace(current, new_name, -1)
            logging.debug("new hosts:\n%s\n", patched)
            fp.write(patched)

        os.system("hostname '%s'" % new_name)
        reboot()


def name():
    global _name
    if _name is None:
        with open('/etc/hostname', 'rt') as fp:
            _name = fp.read().strip()
    return _name


def uptime():
    with open('/proc/uptime') as fp:
        return int(fp.read().split('.')[0])


def mem_usage():
    with open('/proc/meminfo') as fp:
        for line in fp:
            line = line.strip()
            if line.startswith("MemTotal:"):
                kb_mem_total = int(line.split()[1])
            if line.startswith("MemFree:"):
                kb_mem_free = int(line.split()[1])
            if line.startswith("Buffers:"):
                kb_main_buffers = int(line.split()[1])
            if line.startswith("Cached:"):
                kb_main_cached = int(line.split()[1])
        kb_mem_used = kb_mem_total - kb_mem_free - kb_main_cached - kb_main_buffers
        return round(kb_mem_used / kb_mem_total, 1)


def _cpu_stat():
    """
    Returns the split first line of the /proc/stat file
    """
    with open('/proc/stat', 'rt') as fp:
        return list(map(int, fp.readline().split()[1:]))


def cpu_load(tag=None):
    """
    Returns the current cpuload
    """
    if tag and tag in _cpu_stats.keys():
        parts0 = _cpu_stats[tag]
    else:
        parts0 = _cpu_stat()
        time.sleep(0.1)     # only need to sleep when no tag
    parts1 = _cpu_stat()
    if tag:
        _cpu_stats[tag] = parts1

    parts_diff = [p1 - p0 for (p0, p1) in zip(parts0, parts1)]
    user, nice, sys, idle, iowait, irq, softirq, steal, _guest, _guest_nice = parts_diff
    idle_sum = idle + iowait
    non_idle_sum = user + nice + sys + irq + softirq + steal
    total = idle_sum + non_idle_sum
    return non_idle_sum / total


def temperature(celsius=True):
    with open('/sys/class/thermal/thermal_zone0/temp', 'rt') as fp:
        temp = int(fp.read().strip())
    c = int(temp / 1000)
    return c if celsius else ((c * (9 / 5)) + 32)


def shutdown():
    logging.warning("shutting down ...")

    from pwnagotchi.ui import view
    if view.ROOT:
        view.ROOT.on_shutdown()
        # give it some time to refresh the ui
        time.sleep(10)

    logging.warning("syncing...")

    from pwnagotchi import fs
    for m in fs.mounts:
        m.sync()
 
    os.system("sync")
    os.system("halt")


def restart(mode):
    logging.warning("restarting in %s mode ...", mode)
    mode = mode.upper()
    if mode == 'AUTO':
        os.system("touch /root/.pwnagotchi-auto")
    else:
        os.system("touch /root/.pwnagotchi-manual")

    os.system("service bettercap restart")
    time.sleep(1)
    os.system("service pwnagotchi restart")


def reboot(mode=None):
    if mode is not None:
        mode = mode.upper()
        logging.warning("rebooting in %s mode ...", mode)
    else:
        logging.warning("rebooting ...")

    from pwnagotchi.ui import view
    if view.ROOT:
        view.ROOT.on_rebooting()
        # give it some time to refresh the ui
        time.sleep(10)

    if mode == 'AUTO':
        os.system("touch /root/.pwnagotchi-auto")
    elif mode == 'MANU':
        os.system("touch /root/.pwnagotchi-manual")

    logging.warning("syncing...")

    from pwnagotchi import fs
    for m in fs.mounts:
        m.sync()

    os.system("sync")
    os.system("shutdown -r now")

def pwnagotchi_cli():
    def do_clear(display):
        logging.info("clearing the display ...")
        display.clear()
        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.mode = 'auto'
        agent.start()

        while True:
            try:
                # recon on all channels
                agent.recon()
                # get nearby access points grouped by channel
                channels = agent.get_access_points_by_channel()
                # for each channel
                for ch, aps in channels:
                    time.sleep(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 each ap on this channel
                    for ap in aps:
                        # send an association frame in order to get for a PMKID
                        agent.associate(ap)
                        # deauth all client stations in order to get a full handshake
                        for sta in ap['clients']:
                            agent.deauth(ap, sta)
                            time.sleep(1)  # delay to not trigger nexmon firmware bugs

                # An interesting effect of this:
                #
                # From Pwnagotchi's perspective, the more new access points
                # and / or client stations nearby, the longer one epoch of
                # its relative time will take ... basically, in Pwnagotchi's universe,
                # Wi-Fi electromagnetic fields affect time like gravitational fields
                # affect ours ... neat ^_^
                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)

    def add_parsers(parser):
        """
        Adds the plugins and google subcommands
        """
        subparsers = parser.add_subparsers()

        # Add parsers from plugins_cmd
        plugins_cmd.add_parsers(subparsers)

        # Add parsers from google_cmd
        google_cmd.add_parsers(subparsers)

    parser = argparse.ArgumentParser(prog="pwnagotchi")
    # pwnagotchi --help
    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('--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.")

    # Jayofelony added these
    parser.add_argument('--wizard', dest="wizard", action="store_true", default=False,
                        help="Interactive installation of your personal configuration.")
    parser.add_argument('--check-update', dest="check_update", action="store_true", default=False,
                        help="Check for updates on Pwnagotchi. And tells current version.")
    parser.add_argument('--donate', dest="donate", action="store_true", default=False,
                        help="How to donate to this project.")

    # pwnagotchi plugins --help
    add_parsers(parser)
    args = parser.parse_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 google_cmd.used_google_cmd(args):
        config = utils.load_config(args)
        log.setup_logging(args, config)
        rc = google_cmd.handle_cmd(args)
        sys.exit(rc)

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

    if args.wizard:
        def is_valid_hostname(hostname):
            if len(hostname) > 255:
                return False
            if hostname[-1] == ".":
                hostname = hostname[:-1]  # strip exactly one dot from the right, if present
            allowed = re.compile("(?!-)[A-Z\d-]{1,63}(?<!-)$", re.IGNORECASE)
            return all(allowed.match(x) for x in hostname.split("."))

        pwn_restore = input("Do you want to restore the previous configuration?\n\n"
                            "[Y/N]: ")
        if pwn_restore in ('y', 'yes'):
            os.system("cp -f /etc/pwnagotchi/config.toml.bak /etc/pwnagotchi/config.toml")
            print("Your previous configuration is restored, and I will restart in 5 seconds.")
            time.sleep(5)
            os.system("service pwnagotchi restart")
        else:
            pwn_check = input("This will create a new configuration file and overwrite your current backup, are you sure?\n\n"
                              "[Y/N]: ")
            if pwn_check.lower() in ('y', 'yes'):
                os.system("mv -f /etc/pwnagotchi/config.toml /etc/pwnagotchi/config.toml.bak")
                with open("/etc/pwnagotchi/config.toml", "a+") as f:
                    f.write("# Do not edit this file if you do not know what you are doing!!!\n\n")
                    # Set pwnagotchi name
                    print("Welcome to the interactive installation of your personal Pwnagotchi configuration!\n"
                          "My name is Jayofelony, how may I call you?\n\n")
                    pwn_name = input("Pwnagotchi name (no spaces): ")
                    if pwn_name == "":
                        pwn_name = "Pwnagotchi"
                        print("I shall go by Pwnagotchi from now on!")
                        pwn_name = f"main.name = \"{pwn_name}\"\n"
                        f.write(pwn_name)
                    else:
                        if is_valid_hostname(pwn_name):
                            print(f"I shall go by {pwn_name} from now on!")
                            pwn_name = f"main.name = \"{pwn_name}\"\n"
                            f.write(pwn_name)
                        else:
                            print("You have chosen an invalid name. Please start over.")
                            exit()
                    pwn_whitelist = input("How many networks do you want to whitelist? "
                                          "We will also ask a MAC for each network?\n"
                                          "Each SSID and BSSID count as 1 network. \n\n"
                                          "Be sure to use digits as your answer.\n\n"
                                          "Amount of networks: ")
                    if int(pwn_whitelist) > 0:
                        f.write("main.whitelist = [\n")
                        for x in range(int(pwn_whitelist)):
                            ssid = input("SSID (Name): ")
                            bssid = input("BSSID (MAC): ")
                            f.write(f"\t\"{ssid}\",\n")
                            f.write(f"\t\"{bssid}\",\n")
                        f.write("]\n")
                    # set bluetooth tether
                    pwn_bluetooth = input("Do you want to enable BT-Tether?\n\n"
                                          "[Y/N] ")
                    if pwn_bluetooth.lower() in ('y', 'yes'):
                        f.write("main.plugins.bt-tether.enabled = true\n\n")
                        pwn_bluetooth_device = input("What device do you use? Android or iOS?\n\n"
                                                     "Device: ")
                        if pwn_bluetooth_device.lower() == "android":
                            f.write("main.plugins.bt-tether.devices.android-phone.enabled = true\n")
                            pwn_bluetooth_mac = input("What is the bluetooth MAC of your device?\n\n"
                                                      "MAC: ")
                            if pwn_bluetooth_mac != "":
                                f.write(f"main.plugins.bt-tether.devices.android-phone.mac = \"{pwn_bluetooth_mac}\"\n")
                        elif pwn_bluetooth_device.lower() == "ios":
                            f.write("main.plugins.bt-tether.devices.ios-phone.enabled = true\n")
                            pwn_bluetooth_mac = input("What is the bluetooth MAC of your device?\n\n"
                                                      "MAC: ")
                            if pwn_bluetooth_mac != "":
                                f.write(f"main.plugins.bt-tether.devices.ios-phone.mac = \"{pwn_bluetooth_mac}\"\n")
                    # set up display settings
                    pwn_display_enabled = input("Do you want to enable a display?\n\n"
                                                "[Y/N]: ")
                    if pwn_display_enabled.lower() in ('y', 'yes'):
                        f.write("ui.display.enabled = true\n")
                        pwn_display_type = input("What display do you use?\n\n"
                                                 "Be sure to check for the correct display type @ \n"
                                                 "https://github.com/jayofelony/pwnagotchi/blob/master/pwnagotchi/utils.py#L240-L501\n\n"
                                                 "Display type: ")
                        if pwn_display_type != "":
                            f.write(f"ui.display.type = \"{pwn_display_type}\"\n")
                        pwn_display_invert = input("Do you want to invert the display colors?\n"
                                                   "N = Black background\n"
                                                   "Y = White background\n\n"
                                                   "[Y/N]: ")
                        if pwn_display_invert.lower() in ('y', 'yes'):
                            f.write("ui.invert = true\n")
                    f.close()
                    if pwn_bluetooth.lower() in ('y', 'yes'):
                        if pwn_bluetooth_device.lower == "android":
                            print("To visit the webui when connected with your phone, visit: http://192.168.44.44:8080\n"
                                  "Your configuration is done, and I will restart in 5 seconds.")
                        elif pwn_bluetooth_device.lower == "ios":
                            print("To visit the webui when connected with your phone, visit: http://172.20.10.6:8080\n"
                                  "Your configuration is done, and I will restart in 5 seconds.")
                    else:
                        print("Your configuration is done, and I will restart in 5 seconds.")
                    time.sleep(5)
                    os.system("service pwnagotchi restart")
            else:
                print("Ok, doing nothing.")
        sys.exit(0)

    if args.donate:
        print("Donations can be made @ \n "
              "https://github.com/sponsors/jayofelony \n\n"
              "But only if you really want to!")
        sys.exit(0)

    if args.check_update:
        resp = requests.get("https://api.github.com/repos/jayofelony/pwnagotchi/releases/latest")
        latest = resp.json()
        latest_ver = latest['tag_name'].replace('v', '')

        local = version_to_tuple(pwnagotchi.__version__)
        remote = version_to_tuple(latest_ver)
        if remote > local:
            user_input = input("There is a new version available! Update from v%s to v%s?\n[Y/N] " % (pwnagotchi.__version__, latest_ver))
            # input validation
            if user_input.lower() in ('y', 'yes'):
                if os.path.exists('/root/.auto-update'):
                    os.system("rm /root/.auto-update && systemctl restart pwnagotchi")
                else:
                    logging.error("You should make sure auto-update is enabled!")
                print("Okay, give me a couple minutes. Just watch pwnlog while you wait.")
            elif user_input.lower() in ('n', 'no'):  # using this elif for readability
                print("Okay, guess not!")
        else:
            print("You are currently on the latest release, v%s." % 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)
    fonts.init(config)

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

    plugins.load(config)

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

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

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

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

    signal.signal(signal.SIGUSR1, usr1_handler)

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