#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <sqlite3.h>
#include <unistd.h>
#include <crypt.h>

#define DATABASE "/dgldir/dgamelaunch.db"
#define IBUFSIZE 200
#define RANDOMSRC "/dev/urandom"

/* General idea for return status:
 * 0: success
 * 1: password check failed
 * 2: username not found
 * 3: database error
 * 4: invalid input
 */

/* Opens the database and, less obviously, initializes a statement. */
int opendb(sqlite3 **dbp, sqlite3_stmt **stmtp, char *query) {
  int status;
  status = sqlite3_open(DATABASE, dbp);
  if (status) {
    sqlite3_close(*dbp);
    exit(3);
  }
  sqlite3_prepare_v2(*dbp, query, -1, stmtp, NULL);
  if (*stmtp == NULL) {
    sqlite3_close(*dbp);
    exit(3);
  }
  return 0;
}

void cleanup(sqlite3 *db, sqlite3_stmt *stmt, int status) {
  if (stmt)
    sqlite3_finalize(stmt);
  sqlite3_close(db);
  if (status)
    exit(status);
  return;
}

char encode64(unsigned char c) {
  if (c < 26)
    return 'a' + c;
  else if (c < 52)
    return 'A' + c - 26;
  else if (c < 62)
    return '0' + c - 52;
  else if (c == 62)
    return '.';
  else
    return '/';
}

/* Initializes a SHA-512 salt.  salt must contain at least 20 bytes. */
void setsalt(char *salt) {
  unsigned char rnbytes[3], rnvals[4];
  FILE *urandom;
  int loop;
  salt[0] = salt[2] = salt[19] = '$'; /* Delimiters */
  salt[1] = '6'; /* SHA-512 */
  urandom = fopen(RANDOMSRC, "r");
  for (loop = 0; loop < 4; loop++) {
    fread(rnbytes, 1, 3, urandom);
    rnvals[0] = rnbytes[0] >> 2;
    rnvals[1] = ((rnbytes[0] & 0x03) << 4) | ((rnbytes[1] & 0xf0) >> 4);
    rnvals[2] = ((rnbytes[1] & 0x0f) << 2) | ((rnbytes[2] & 0xc0) >> 6);
    rnvals[3] = rnbytes[2] & 0x3f;
    salt[loop * 4 + 3] = encode64(rnvals[0]);
    salt[loop * 4 + 4] = encode64(rnvals[1]);
    salt[loop * 4 + 5] = encode64(rnvals[2]);
    salt[loop * 4 + 6] = encode64(rnvals[3]);
  }
  fclose(urandom);
}

int check(char *uname, char *pw) {
  char *pwhash, *comphash;
  char *query = "SELECT password FROM dglusers WHERE username=?;";
  int status;
  sqlite3 *db;
  sqlite3_stmt *qstmt;

  opendb(&db, &qstmt, query);
  status = sqlite3_bind_text(qstmt, 1, uname, -1, SQLITE_TRANSIENT);
  if (status)
    cleanup(db, qstmt, 3);
  status = sqlite3_step(qstmt);
  if (status != SQLITE_ROW) {
    sqlite3_finalize(qstmt);
    sqlite3_close(db);
    if (status == SQLITE_DONE)
      return 2; /* User not found */
    return 3;
  }
  pwhash = strdup((char *) sqlite3_column_text(qstmt, 0));
  cleanup(db, qstmt, 0);

  /* Check the password */
  comphash = crypt(pw, pwhash);
  if (!strcmp(pwhash, comphash))
    status = 0;
  else
    status = 1;
  free(pwhash);
  return status;
}

int insertuser(char *uname, char *pw, char *email) {
  char *checkquery = "SELECT * FROM dglusers WHERE username = ?;";
  char *addquery =  "INSERT INTO dglusers (username, password, email, flags, env) VALUES (?, ?, ?, ?, ?);";
  int status;
  sqlite3 *db;
  sqlite3_stmt *qstmt;
  char salt[20];

  opendb(&db, &qstmt, checkquery);
  /* Check for existing account in the same transaction with creating it. */
  status = sqlite3_exec(db, "BEGIN;", NULL, NULL, NULL);
  if (status)
    cleanup(db, qstmt, 3);
  sqlite3_bind_text(qstmt, 1, uname, -1, SQLITE_TRANSIENT);
  status = sqlite3_step(qstmt);
  if (status != SQLITE_DONE) {
    sqlite3_finalize(qstmt);
    sqlite3_close(db);
    if (status == SQLITE_ROW)
      return 1; /* User already exists */
    return 3;
  }
  /* The username doesn't exist yet, so create a new account. */
  setsalt(salt);
  sqlite3_finalize(qstmt);
  sqlite3_prepare_v2(db, addquery, -1, &qstmt, NULL);
  if (qstmt == NULL)
    cleanup(db, NULL, 3);
  sqlite3_bind_text(qstmt, 1, uname, -1, SQLITE_TRANSIENT);
  sqlite3_bind_text(qstmt, 2, strdup(crypt(pw, salt)), -1, free);
  sqlite3_bind_text(qstmt, 3, email, -1, SQLITE_TRANSIENT);
  sqlite3_bind_int(qstmt, 4, 0);
  sqlite3_bind_text(qstmt, 5, "", -1, SQLITE_STATIC);
  status = sqlite3_step(qstmt);
  if (status != SQLITE_DONE)
    cleanup(db, qstmt, 3);
  status = sqlite3_exec(db, "COMMIT;", NULL, NULL, NULL);
  cleanup(db, qstmt, 0);
  return status;
}

int getmail(char *uname) {
  char *mailquery = "SELECT email FROM dglusers WHERE username = ?;";
  int status;
  sqlite3 *db;
  sqlite3_stmt *qstmt;

  opendb(&db, &qstmt, mailquery);
  status = sqlite3_bind_text(qstmt, 1, uname, -1, SQLITE_TRANSIENT);
  if (status)
    cleanup(db, qstmt, 3);
  status = sqlite3_step(qstmt);
  if (status != SQLITE_ROW) {
    sqlite3_finalize(qstmt);
    sqlite3_close(db);
    if (status == SQLITE_DONE)
      return 2; /* User not found */
    return 3;
  }
  printf("%s\n", sqlite3_column_text(qstmt, 0));
  cleanup(db, qstmt, 0);
  return 0;
}

int setmail(char *uname, char *newmail) {
  char *setquery = "UPDATE dglusers SET email = ? WHERE username = ?;";
  sqlite3 *db;
  sqlite3_stmt *qstmt;
  int status;

  opendb(&db, &qstmt, setquery);
  status = sqlite3_bind_text(qstmt, 2, uname, -1, SQLITE_TRANSIENT);
  if (status)
    cleanup(db, qstmt, 3);
  status = sqlite3_bind_text(qstmt, 1, newmail, -1, SQLITE_TRANSIENT);
  if (status)
    cleanup(db, qstmt, 3);
  status = sqlite3_step(qstmt);
  if (status != SQLITE_DONE)
    cleanup(db, qstmt, 3);
  cleanup(db, qstmt, 0);
  return 0;
}

int setpw(char *uname, char *newpw) {
  char *setquery = "UPDATE dglusers SET password = ? WHERE username = ?;";
  sqlite3 *db;
  sqlite3_stmt *qstmt;
  int status;
  char *hash;
  char salt[20];

  setsalt(salt);
  hash = crypt(newpw, salt);
  opendb(&db, &qstmt, setquery);
  status = sqlite3_bind_text(qstmt, 2, uname, -1, SQLITE_TRANSIENT);
  if (status)
    cleanup(db, qstmt, 3);
  status = sqlite3_bind_text(qstmt, 1, hash, -1, SQLITE_TRANSIENT);
  if (status)
    cleanup(db, qstmt, 3);
  status = sqlite3_step(qstmt);
  if (status != SQLITE_DONE)
    cleanup(db, qstmt, 3);
  cleanup(db, qstmt, 0);
  return 0;
}

int main(int argc, char *argv[]) {
  char ibuf[IBUFSIZE], *uname, *line1, *line2;
  char *cptr; // Utility pointer
  int status;

  /* Read in the username and password */
  fgets(ibuf, IBUFSIZE, stdin);
  uname = ibuf;
  line1 = strchr(uname, '\n');
  if (line1 == NULL)
    exit(4); /* Truncated */
  *line1 = '\0';
  if (argc == 1 || strcmp(argv[1], "getmail")) {
    line1++;
    fgets(line1, IBUFSIZE - (line1 - ibuf), stdin);
    if (line1[strlen(line1) - 1] == '\n')
      line1[strlen(line1) - 1] = '\0';
    else
      exit(4); /* Truncated */
  }
  if (argc > 1 && !strcmp(argv[1], "setmail")) {
    /* E-mail, sanitize */
    for (cptr = line1; *cptr != '\0'; cptr++) {
      if (!isalnum(*cptr) && !strchr("@._-", *cptr)) {
        exit(4);
      }
    }
  }
  if (argc > 1 && !strcmp(argv[1], "register")) {
    line2 = line1 + strlen(line1) + 1;
    fgets(line2, IBUFSIZE - (line2 - ibuf), stdin);
    if (line2[strlen(line2) - 1] == '\n')
      line2[strlen(line2) - 1] = '\0';
    else
      exit(4);
    for (cptr = line2; *cptr != '\0'; cptr++) {
      if (!isalnum(*cptr) && !strchr("@._-", *cptr)) {
        exit(4);
      }
    }
  }
  /* Sanitize the username, because it gets put into a query. */
  for (cptr = uname; *cptr != '\0'; cptr++) {
    if (!isalnum(*cptr)) {
      exit(4);
    }
  }
  if (argc == 1 || !strcmp(argv[1], "check"))
    status = check(uname, line1);
  else if (!strcmp(argv[1], "register")) {
    status = insertuser(uname, line1, line2);
  }
  else if (!strcmp(argv[1], "getmail")) {
    status = getmail(uname);
  }
  else if (!strcmp(argv[1], "setmail")) {
    status = setmail(uname, line1);
  }
  else if (!strcmp(argv[1], "setpw")) {
    status = setpw(uname, line1);
  }
  else
    status = 127;
  return status;
}