/*
 * wpa_actiond.c
 * Copyright (c) 2009, Thomas Bächler <thomas@archlinux.org>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 *
 * See COPYING for more details.
 */

#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <signal.h>
#include <sys/select.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/un.h>
#include <malloc.h>
#include <syslog.h>
#include "wpa_ctrl/wpa_ctrl.h"

/* buffer sizes */
#define BUFLEN 1024
#define SSIDLEN 256
#define IDSTRLEN 256
/* timeout (in seconds) after which wpa_actiond considers a lost
   wlan connection "disconnected" */
#define DEFAULT_DISCONNECT_TIMEOUT 30
/* default paths */
#define DEFAULT_CTRL_PATH "/var/run/wpa_supplicant"
#define DEFAULT_PIDFILE "/var/run/wpa_actiond.pid"

enum wpaevent {
  WPA_ACTIOND_EVENT_UNCHANGED,
  WPA_ACTIOND_EVENT_CONNECTED,
  WPA_ACTIOND_EVENT_DISCONNECTED,
  WPA_ACTIOND_EVENT_TERMINATED
};

enum wpastate {
  WPA_ACTIOND_STATE_UNDEFINED,
  WPA_ACTIOND_STATE_CONNECTED,
  WPA_ACTIOND_STATE_DISCONNECTED,
  WPA_ACTIOND_STATE_CONNECTION_LOST,
  WPA_ACTIOND_STATE_TERMINATED
};

enum wpaaction {
  WPA_ACTIOND_ACTION_CONNECT,
  WPA_ACTIOND_ACTION_CONNECTION_LOST,
  WPA_ACTIOND_ACTION_CONNECTION_REESTABLISHED,
  WPA_ACTIOND_ACTION_DISCONNECT
};

enum wpa_actiond_logevent {
  WPA_ACTIOND_LOG_SET_STDERR,
  WPA_ACTIOND_LOG_STARTED,
  WPA_ACTIOND_LOG_CONNECTED,
  WPA_ACTIOND_LOG_DISCONNECTED,
  WPA_ACTIOND_LOG_CONNECTION_LOST,
  WPA_ACTIOND_LOG_CONNECTION_REESTABLISHED,
  WPA_ACTIOND_LOG_TERMINATE,
  WPA_ACTIOND_LOG_ERROR,
  WPA_ACTIOND_LOG_CUSTOM_ERROR
};

static void logevent(enum wpa_actiond_logevent l, const char *iface, const char *arg) {
  static int isopen = 0, tostderr = 0;

  if(l == WPA_ACTIOND_LOG_SET_STDERR) {
    tostderr = 1;
    return;
  }

  if(!isopen && !tostderr) {
    openlog("wpa_actiond", LOG_PID, LOG_DAEMON);
    isopen = 1;
  }

  switch(l) {
    case WPA_ACTIOND_LOG_STARTED:
      if(!tostderr)
        syslog(LOG_NOTICE, "Starting wpa_actiond session for interface '%s'", iface);
      else
        fprintf(stderr, "Starting wpa_actiond session for interface '%s'\n", iface);
      break;
    case WPA_ACTIOND_LOG_CONNECTED:
      if(!tostderr)
        syslog(LOG_NOTICE, "Interface '%s' connected to network '%s'", iface, arg);
      else
        fprintf(stderr, "Interface '%s' connected to network '%s'\n", iface, arg);
      break;
    case WPA_ACTIOND_LOG_DISCONNECTED:
      if(!tostderr)
        syslog(LOG_NOTICE, "Interface '%s' disconnected from network '%s'", iface, arg);
      else
        fprintf(stderr, "Interface '%s' disconnected from network '%s'\n", iface, arg);
      break;
    case WPA_ACTIOND_LOG_CONNECTION_LOST:
      if(!tostderr)
        syslog(LOG_NOTICE, "Interface '%s' lost connection to network '%s'", iface, arg);
      else
        fprintf(stderr, "Interface '%s' lost connection to network '%s'\n", iface, arg);
      break;
    case WPA_ACTIOND_LOG_CONNECTION_REESTABLISHED:
      if(!tostderr)
        syslog(LOG_NOTICE, "Interface '%s' reestablished connection to network '%s'", iface, arg);
      else
        fprintf(stderr, "Interface '%s' reestablished connection to network '%s'\n", iface, arg);
      break;
    case WPA_ACTIOND_LOG_TERMINATE:
      if(!tostderr)
        syslog(LOG_NOTICE, "Terminating wpa_actiond session for interface '%s'", iface);
      else
        fprintf(stderr, "Terminating wpa_actiond session for interface '%s'\n", iface);
      break;
    case WPA_ACTIOND_LOG_ERROR:
      if(!tostderr)
        syslog(LOG_ERR, "Error (%s): %s%m", iface, arg);
      else
        fprintf(stderr, "Error (%s): %s%s\n", iface, arg, strerror(errno));
      break;
    case WPA_ACTIOND_LOG_CUSTOM_ERROR:
      if(!tostderr)
        syslog(LOG_ERR, "Error (%s): %s", iface, arg);
      else
        fprintf(stderr, "Error (%s): %s\n", iface, arg);
      break;
  }
}

static int min(int i, int j) {
  return i < j ? i : j;
}

static void get_ssid(struct wpa_ctrl *ctrl, char *ssid) {
  char reply[BUFLEN], *pos = &reply[0], *pos2;
  size_t reply_len = BUFLEN-1;
  *ssid = '\0';

  if(wpa_ctrl_request(ctrl, "STATUS", strlen("STATUS"), reply, &reply_len, NULL))
    return;
  reply[reply_len] = '\0';

  while(1) {
    pos = strstr(pos, "ssid=");
    if(pos == NULL || pos == reply || *(pos-1) == '\n')
      break;
    pos += strlen("ssid=");
  }

  if(pos != NULL) {
    pos += strlen("ssid=");
    pos2 = index(pos, '\n');
    strncpy(ssid, pos, min(pos2-pos, SSIDLEN-1));
    ssid[min(pos2-pos, SSIDLEN-1)] = '\0';
  }
}

static int str_match(const char *s1, const char *s2) {
  return !strncmp(s1,s2,strlen(s2));
}

static void parse_wpa_event(const char *msg, enum wpaevent *ev, char *idstr)
{
  const char *pos, *pos2;

  *ev = WPA_ACTIOND_EVENT_UNCHANGED;

  pos = msg;
  if (*pos == '<') {
    /* skip priority */
    pos = strchr(pos, '>');
    if(pos)
      pos++;
    else
      pos = msg;
  }

  if (str_match(pos, WPA_EVENT_CONNECTED)) {
    *ev = WPA_ACTIOND_EVENT_CONNECTED;
    pos = msg;
    pos = strstr(pos, "id_str=");
    if(pos != NULL) {
      pos += strlen("id_str=");
      pos2 = rindex(pos, ']');
      strncpy(idstr, pos, min(pos2-pos, IDSTRLEN-1));
      idstr[min(pos2-pos, IDSTRLEN-1)] = '\0';
    }
  }
  else if (str_match(pos, WPA_EVENT_DISCONNECTED))
    *ev = WPA_ACTIOND_EVENT_DISCONNECTED;
  else if (str_match(pos, WPA_EVENT_TERMINATING))
    *ev = WPA_ACTIOND_EVENT_TERMINATED;
}

static void action(enum wpaaction act, const char *iface, const char *ssid, const char *idstr, const int ctrlfd, const char *script) {
  char *actstr = NULL;
  pid_t f;

  switch(act) {
    case WPA_ACTIOND_ACTION_CONNECT:
      actstr = "CONNECT";
      break;
    case WPA_ACTIOND_ACTION_CONNECTION_LOST:
      actstr = "LOST";
      break;
    case WPA_ACTIOND_ACTION_CONNECTION_REESTABLISHED:
      actstr = "REESTABLISHED";
      break;
    case WPA_ACTIOND_ACTION_DISCONNECT:
      actstr = "DISCONNECT";
      break;
  }

  if(actstr == NULL)
    return;

  f = fork();
  switch(f) {
    case -1:
      logevent(WPA_ACTIOND_LOG_ERROR, iface, "fork(): ");
      break;
    case 0:
      close(ctrlfd);
      signal(SIGTERM, SIG_DFL);
      signal(SIGHUP, SIG_DFL);
      execl(script, script, iface, ssid, idstr, actstr, (char*)NULL);
      logevent(WPA_ACTIOND_LOG_ERROR, iface, "execl(): ");
      exit(-1);
      break;
    default:
      waitpid(f, NULL, 0);
      break;
  }
}

static void loop(const char *iface, const char *ctrlpath, const int disconnect_timeout, const char *script) {
  /* wpa_supplicant control structure */
  struct wpa_ctrl *ctrl;
  /* buffer for wpa_supplicant replies */
  char reply[BUFLEN];
  size_t reply_len;
  /* states and events */
  enum wpastate state = WPA_ACTIOND_STATE_UNDEFINED;
  enum wpaevent ev;
  /* select stuff */
  int ctrl_fd;
  fd_set ctrl_fds;
  int r;
  /* select timeout */
  struct timeval timeout;
  /* save ssid */
  char ssid[SSIDLEN], old_ssid[SSIDLEN];
  char idstr[IDSTRLEN], old_idstr[IDSTRLEN];
  /* path to control socket */
  char *ctrlsock = NULL;
  int ctrlsocklen;

  ssid[0] = '\0';
  old_ssid[0] = '\0';
  idstr[0] = '\0';
  old_idstr[0] = '\0';

  /* set up signals */
  void terminate(int s) {
    if(state == WPA_ACTIOND_STATE_CONNECTED || state == WPA_ACTIOND_STATE_CONNECTION_LOST) {
      logevent(WPA_ACTIOND_LOG_DISCONNECTED, iface, ssid);
      action(WPA_ACTIOND_ACTION_DISCONNECT, iface, ssid, idstr, wpa_ctrl_get_fd(ctrl), script);
    }
    logevent(WPA_ACTIOND_LOG_TERMINATE, iface, "");

    FD_ZERO(&ctrl_fds);
    wpa_ctrl_detach(ctrl);
    wpa_ctrl_close(ctrl);
    exit(0);
  }
  signal(SIGTERM, terminate);
  signal(SIGHUP, SIG_IGN);

  /* open wpa_supplicant cotrol interface */
  ctrlsocklen = strlen(iface)+strlen(ctrlpath)+2;
  ctrlsock = (char*)malloc(ctrlsocklen);
  if(ctrlsock == NULL) {
    logevent(WPA_ACTIOND_LOG_ERROR, iface, "malloc(): ");
    exit(-1);
  }
  snprintf(ctrlsock, ctrlsocklen, "%s/%s", ctrlpath, iface);
  ctrl = wpa_ctrl_open(ctrlsock);
  free(ctrlsock);
  if(ctrl == NULL) {
    logevent(WPA_ACTIOND_LOG_CUSTOM_ERROR, iface, "Unable to open wpa_supplicant control socket");
    exit(-1);
  }
  if(wpa_ctrl_attach(ctrl)) {
    logevent(WPA_ACTIOND_LOG_CUSTOM_ERROR, iface, "Could not attach to wpa_supplicant");
    wpa_ctrl_close(ctrl);
    exit(-1);
  }

  /* set initial ssid */
  get_ssid(ctrl, ssid);

  logevent(WPA_ACTIOND_LOG_STARTED, iface, "");

  while(state != WPA_ACTIOND_STATE_TERMINATED) {
    /* these parameters may have been altered */
    FD_ZERO(&ctrl_fds);
    ctrl_fd = wpa_ctrl_get_fd(ctrl);
    FD_SET(ctrl_fd, &ctrl_fds);
    /* wait for an event */
    r = select(ctrl_fd+1, &ctrl_fds, NULL, NULL, (state == WPA_ACTIOND_STATE_CONNECTION_LOST) ? &timeout : NULL);
    switch(r) {
      case -1:
        /* select error, terminate wpa_actiond */
        logevent(WPA_ACTIOND_LOG_ERROR, iface, "select(): ");
        if(state == WPA_ACTIOND_STATE_CONNECTED || state == WPA_ACTIOND_STATE_CONNECTION_LOST) {
          logevent(WPA_ACTIOND_LOG_DISCONNECTED, iface, ssid);
          action(WPA_ACTIOND_ACTION_DISCONNECT, iface, ssid, idstr, wpa_ctrl_get_fd(ctrl), script);
        }
        logevent(WPA_ACTIOND_LOG_TERMINATE, iface, "");
        state = WPA_ACTIOND_STATE_TERMINATED;
        break;
      case 0:
        /* timeout, we are in WPA_ACTIOND_STATE_CONNECTION_LOST, disconnect */
        logevent(WPA_ACTIOND_LOG_DISCONNECTED, iface, ssid);
        state = WPA_ACTIOND_STATE_DISCONNECTED;
        action(WPA_ACTIOND_ACTION_DISCONNECT, iface, ssid, idstr, wpa_ctrl_get_fd(ctrl), script);
        break;
      default:
        /* event received */
        reply_len = BUFLEN-1;
        wpa_ctrl_recv(ctrl, reply, &reply_len);
        reply[reply_len] = '\0';
        parse_wpa_event(reply, &ev, idstr);
        switch(ev) {
          case WPA_ACTIOND_EVENT_CONNECTED:
            /* connect to wireless */
            get_ssid(ctrl, ssid);
            if(state == WPA_ACTIOND_STATE_CONNECTION_LOST) {
              if(strcmp(old_ssid, ssid)) {
                logevent(WPA_ACTIOND_LOG_DISCONNECTED, iface, old_ssid);
                action(WPA_ACTIOND_ACTION_DISCONNECT, iface, old_ssid, old_idstr, wpa_ctrl_get_fd(ctrl), script);
                logevent(WPA_ACTIOND_LOG_CONNECTED, iface, ssid);
                action(WPA_ACTIOND_ACTION_CONNECT, iface, ssid, idstr, wpa_ctrl_get_fd(ctrl), script);
              }
              else {
                logevent(WPA_ACTIOND_LOG_CONNECTION_REESTABLISHED, iface, ssid);
                action(WPA_ACTIOND_ACTION_CONNECTION_REESTABLISHED, iface, ssid, idstr, wpa_ctrl_get_fd(ctrl), script);
              }
            }
            else {
              logevent(WPA_ACTIOND_LOG_CONNECTED, iface, ssid);
              action(WPA_ACTIOND_ACTION_CONNECT, iface, ssid, idstr, wpa_ctrl_get_fd(ctrl), script);
            }
            state = WPA_ACTIOND_STATE_CONNECTED;
            break;
          case WPA_ACTIOND_EVENT_DISCONNECTED:
            if(state == WPA_ACTIOND_STATE_CONNECTED) {
              /* we lost the connection, save old ssid, wait disconnect_timeout seconds */
              logevent(WPA_ACTIOND_LOG_CONNECTION_LOST, iface, ssid);
              action(WPA_ACTIOND_ACTION_CONNECTION_LOST, iface, ssid, idstr, wpa_ctrl_get_fd(ctrl), script);
              state = WPA_ACTIOND_STATE_CONNECTION_LOST;
              strncpy(old_ssid, ssid, SSIDLEN);
              strncpy(old_idstr, idstr, IDSTRLEN);

              /* Warning: If we are in state CONNECTION_LOST and receive any event
               * other than CONNECT, timeout will not be reset to its initial value.
               * This will work as expected on Linux, but not on other platforms
               */
              timeout.tv_sec = disconnect_timeout;
              timeout.tv_usec = 0;
            }
            break;
          case WPA_ACTIOND_EVENT_TERMINATED:
            /* supplicant has been terminated */
            if(state == WPA_ACTIOND_STATE_CONNECTED || state == WPA_ACTIOND_STATE_CONNECTION_LOST) {
              logevent(WPA_ACTIOND_LOG_DISCONNECTED, iface, ssid);
              action(WPA_ACTIOND_ACTION_DISCONNECT, iface, ssid, idstr, wpa_ctrl_get_fd(ctrl), script);
            }
            logevent(WPA_ACTIOND_LOG_TERMINATE, iface, "");
            state = WPA_ACTIOND_STATE_TERMINATED;
            break;
          case WPA_ACTIOND_EVENT_UNCHANGED:
            /* we are not interested in this event */
            break;
        }
    }
  }

  signal(SIGTERM, SIG_IGN);
  FD_ZERO(&ctrl_fds);
  wpa_ctrl_detach(ctrl);
  wpa_ctrl_close(ctrl);
}

static void usage(const char *progname) {
  fprintf(stderr, "\nUsage:\n"
                  "  %s -i interface -a script [-n] [-p control_socket_path] [-P pid_file] [-t timeout]\n\n"
                  "  -i\tname of the wireless interface\n"
                  "  -a\taction script\n"
                  "  -n\tdo not daemonize, log to stderr (default: daemonize, log to syslog)\n"
                  "  -p\tpath to the wpa_supplicant control socket directory (default: %s)\n"
                  "  -P\twpa_actiond PID file (default: %s)\n"
                  "  -t\ttimeout in seconds until a lost connection is considered disconnected (default: %d)\n\n"
                  "  The action script takes four parameters:\n"
                  "    1) The interface name\n"
                  "    2) The SSID of the network\n"
                  "    3) The id_str parameter of the wpa_supplicant network section\n"
                  "    4) One of the strings CONNECT, LOST, REESTABLISHED and DISCONNECT\n"
                  "\n", progname, DEFAULT_CTRL_PATH, DEFAULT_PIDFILE, DEFAULT_DISCONNECT_TIMEOUT);
}

int main(int argc, char *argv[]) {
  int opt, daemonize = 1, timeout = DEFAULT_DISCONNECT_TIMEOUT;
  const char *iface = NULL, *ctrlpath = NULL, *pidfile = NULL, *script = NULL;
  FILE *pid;

  /* read options */
  while((opt = getopt(argc,argv,"ni:p:P:t:a:")) != -1) {
    switch(opt) {
      case 'n':
        daemonize = 0;
        break;
      case 'i':
        iface = optarg;
        break;
      case 'p':
        ctrlpath = optarg;
        break;
      case 'P':
        pidfile = optarg;
        break;
      case 't':
        timeout = atoi(optarg);
        if(timeout<=0) {
          usage(argv[0]);
          return -1;
        }
        break;
      case 'a':
        script = optarg;
        break;
      default:
        usage(argv[0]);
        return -1;
    }
  }
  if(iface == NULL || script == NULL) {
    usage(argv[0]);
    return -1;
  }
  /* set defaults for unset options */
  if(ctrlpath == NULL)
    ctrlpath = DEFAULT_CTRL_PATH;
  if(pidfile == NULL)
    pidfile = DEFAULT_PIDFILE;

  if(daemonize) {
    /* go to background */
    if(daemon(0,0)) {
      logevent(WPA_ACTIOND_LOG_ERROR, iface, "daemon(): ");
      return -1;
    }
  }
  else {
    /* redirect logging to stderr instead of syslog */
    logevent(WPA_ACTIOND_LOG_SET_STDERR, NULL, NULL);
  }
  /* write pid file */
  pid = fopen(pidfile, "w");
  if(pid == NULL) {
    logevent(WPA_ACTIOND_LOG_ERROR, iface, "fopen(): ");
    return -1;
  }
  fprintf(pid, "%d\n", getpid());
  fclose(pid);
  loop(iface, ctrlpath, timeout, script);
  return 0;
}
