twofap/mainwindow.cpp

#include <QHBoxLayout>
#include <QVBoxLayout>
#include <QToolButton>
#include <QAction>
#include <QSpacerItem>
#include <QInputDialog>
#include <QRegularExpression>
#include <QMessageBox>
#include <QSettings>
#include <QCoreApplication>
#include <QTimer>
#include <QScrollArea>
#include <QLabel>
#include "mainwindow.h"
#include "libcppotp/bytes.h"
#include "libcppotp/otp.h"

extern "C" {
#include <time.h>
}

#define GROUP_NAME "keys"
#define TIMER_INTERVAL_MS 1000

MainWindow::MainWindow(QWidget *parent) :
  QMainWindow(parent)
  ,m_centralWidget(nullptr)
  ,m_settings(qApp->applicationName() + ".ini", QSettings::IniFormat)
  ,m_statusLabel(nullptr)
  ,m_timer(new QTimer())
{
  m_settings.setIniCodec("UTF-8");

  if(!m_settings.isWritable()) {
    QMessageBox::critical(0, tr("Error"), tr("Permission denied while trying to write to settings file."));
    qApp->quit();
    return;
  }

  QByteArray savedGeometry = m_settings.value("main/geometry").toByteArray();

  if(savedGeometry.size() > 0)
    restoreGeometry(savedGeometry);

  QTimer::singleShot(0, this, &MainWindow::rebuildLayout);

  connect(this, &MainWindow::addKeyClicked, this, &MainWindow::onAddKeyClicked);

  connect(m_timer, &QTimer::timeout, this, [this]() {
    if(!m_statusLabel)
      return;

    unsigned secs = 30 - (time(NULL) % 30);
    QString remaining = QString::number(secs) + " seconds remaining.";

    m_statusLabel->setText(remaining);

    if(secs == 30)
      QTimer::singleShot(0, this, &MainWindow::rebuildLayout);
  });
}

MainWindow::~MainWindow() {
  m_settings.setValue("main/geometry", saveGeometry());
  m_settings.sync();
}

void MainWindow::rebuildLayout() {
  if(m_timer)
    m_timer->stop();

  if(m_statusLabel)
    delete m_statusLabel;

  if(m_centralWidget)
    delete m_centralWidget;

  // default layout is a vertical box
  m_centralWidget = new QWidget();
  setCentralWidget(m_centralWidget);

  QVBoxLayout *vbox = new QVBoxLayout();
  m_centralWidget->setLayout(vbox);

  // main vbox layout for codes
  QWidget *codeWidget = new QWidget();
  QVBoxLayout *codeVBoxLayout = new QVBoxLayout();

  codeWidget->setLayout(codeVBoxLayout);

  QScrollArea *scrollArea = new QScrollArea();
  scrollArea->setWidget(codeWidget);
  scrollArea->setWidgetResizable(true);

  vbox->addWidget(scrollArea);

  // generate each code for keys in settings file and add to layout
  m_settings.beginGroup(GROUP_NAME);

  QStringList names = m_settings.childKeys();
  names.sort(Qt::CaseInsensitive);

  for(int i = 0; i < names.size(); ++i) {
    QString name = names.at(i);
    QString key = m_settings.value(name).toString();
    QString code = generateCode(key);
    
    QLabel *label = new QLabel(name + ": " + code, 0);
    codeVBoxLayout->addWidget(label);
  }

  m_settings.endGroup();

  // status bar at the bottom
  QHBoxLayout *statusBoxLayout = new QHBoxLayout();

  QToolButton *addButton = new QToolButton();

  QAction *addAction = new QAction(tr("Add"));
  connect(addAction, &QAction::triggered, this, &MainWindow::addKeyClicked);

  addButton->setDefaultAction(addAction);

  statusBoxLayout->addWidget(addButton);
  statusBoxLayout->addSpacerItem(new QSpacerItem(20, 20, QSizePolicy::Expanding, QSizePolicy::Minimum));

  m_statusLabel = new QLabel();

  statusBoxLayout->addWidget(m_statusLabel);

  vbox->addSpacerItem(new QSpacerItem(20, 20, QSizePolicy::Minimum, QSizePolicy::Expanding));
  vbox->addLayout(statusBoxLayout);

  m_timer->start(TIMER_INTERVAL_MS);
}

void MainWindow::onAddKeyClicked() {
  // get name and key for this account and save to settings
  QString name = QInputDialog::getText(this, tr("Add Key"), tr("Please enter a name for this account:"));

  if(name.isEmpty()) {
    QMessageBox::critical(this, tr("Error"), tr("No name entered."));
    return;
  }

  QString key = QInputDialog::getText(this, tr("Add Key"), tr("Please enter the TOTP key for this account:"));

  // sanitize the key before saving (remove any spaces or invalid characters)
  key = key.replace(QRegularExpression(R"([^a-zA-Z2-7])"), "");
  key = key.toUpper();

  if(key.isEmpty()) {
    QMessageBox::critical(this, tr("Error"), tr("No key entered."));
    return;
  }

  saveNewKey(name, key);

  // a bit lazy and heavy-handed, but works for now
  rebuildLayout();
}

void MainWindow::saveNewKey(QString name, QString key) {
  // newly added key goes into our ini file
  m_settings.setValue(QString(GROUP_NAME) + "/" + name, key);
  m_settings.sync();
}

QString MainWindow::generateCode(QString key) {
  // where the magic happens
  uint32_t p = 0;
  CppTotp::Bytes::ByteString qui;

  try {
    std::string ckey = key.toStdString();
    qui = CppTotp::Bytes::fromUnpaddedBase32(ckey);
  } catch(const std::invalid_argument&) {
    QMessageBox::critical(0, tr("Error"), tr("Invalid key encountered, exiting."));
    qApp->quit();
    return QString();
  }

  p = CppTotp::totp(qui, time(NULL), 0, 30, 6);

  // return only the first six digits, as a string, with a space in the middle
  return QString::number(p).left(6).replace(QRegularExpression(R"((\d{3})(\d{3}))"), R"(\1 \2)");
}