main.cpp

#include <QApplication>
#include <QCheckBox>
#include <QDialog>
#include <QDialogButtonBox>
#include <QDir>
#include <QDirIterator>
#include <QFile>
#include <QFileInfo>
#include <QFormLayout>
#include <QGroupBox>
#include <QHeaderView>
#include <QLabel>
#include <QMessageBox>
#include <QProcess>
#include <QPushButton>
#include <QTemporaryFile>
#include <QScrollArea>
#include <QSet>
#include <QTableWidget>
#include <QTableWidgetItem>
#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QDateTime>
#include <algorithm>
#include <utility>

namespace {

    struct WakeDevice {
        QString wakeFile;
        QString sysfsPath;
        QString name;
        QString kind;
        bool enabled = false;
        bool pending = false;
    };

    QString readText(const QString &path)
    {
        QFile file(path);
        if (!file.open(QIODevice::ReadOnly))
            return {};
        return QString::fromLocal8Bit(file.readAll()).trimmed();
    }

    QString attributeOrEmpty(const QString &path, const QString &attribute)
    {
        return readText(path + QLatin1Char('/') + attribute);
    }

    QString friendlyName(const QString &path)
    {
        const QString product = attributeOrEmpty(path, QStringLiteral("product"));
        const QString name = attributeOrEmpty(path, QStringLiteral("name"));
        const QString modalias = attributeOrEmpty(path, QStringLiteral("modalias"));
        const QString fileName = QFileInfo(path).fileName();

        // Common laptop wake sources deserve stable, human-readable names.
        if (path.contains(QStringLiteral("/serio")))
            return QStringLiteral("Built-in keyboard");
        if (path.contains(QStringLiteral("gpio-keys")))
            return QStringLiteral("Power button (GPIO)");
        if (modalias.contains(QStringLiteral("PNP0C0D")))
            return QStringLiteral("Lid switch");
        if (modalias.contains(QStringLiteral("ELAN")) || path.contains(QStringLiteral("i2c-ELAN")))
            return QStringLiteral("ELAN touchpad");
        if (modalias.contains(QStringLiteral("PNP0B00")))
            return QStringLiteral("RTC clock");
        if (path.contains(QStringLiteral("alarmtimer")))
            return QStringLiteral("RTC alarm timer");
        if (path.contains(QStringLiteral("power_supply")))
            return product.isEmpty() ? QStringLiteral("Power supply (%1)").arg(fileName)
            : QStringLiteral("Power supply: %1").arg(product);

        if (path.contains(QStringLiteral("/usb"))) {
            const QString manufacturer = attributeOrEmpty(path, QStringLiteral("manufacturer"));
            if (!product.isEmpty() && !manufacturer.isEmpty())
                return QStringLiteral("%1 %2").arg(manufacturer, product);
            if (!product.isEmpty())
                return product;
            if (!manufacturer.isEmpty())
                return manufacturer + QStringLiteral(" USB device");
            return QStringLiteral("USB device (%1)").arg(fileName);
        }

        if (modalias.startsWith(QStringLiteral("pci:"))) {
            if (path.contains(QStringLiteral("pcieport")))
                return QStringLiteral("PCIe root port");
            if (modalias.contains(QStringLiteral("bc0Csc03")))
                return QStringLiteral("USB controller");
            if (modalias.contains(QStringLiteral("bc04")))
                return QStringLiteral("Audio controller");
            return QStringLiteral("PCI device (%1)").arg(fileName);
        }

        if (!product.isEmpty())
            return product;
        if (!name.isEmpty())
            return name;
        if (modalias.startsWith(QStringLiteral("acpi:")))
            return QStringLiteral("ACPI device (%1)").arg(fileName);
        if (modalias.startsWith(QStringLiteral("platform:")))
            return QStringLiteral("Platform device (%1)").arg(fileName);
        if (modalias.startsWith(QStringLiteral("serio:")))
            return QStringLiteral("Keyboard controller");
        return fileName.isEmpty() ? QStringLiteral("Unnamed device")
        : QStringLiteral("Device (%1)").arg(fileName);
    }

    QString deviceKind(const QString &path)
    {
        if (path.contains("/usb"))
            return QStringLiteral("USB");
        if (path.contains("/pci"))
            return QStringLiteral("PCI");
        if (path.contains("/i2c"))
            return QStringLiteral("I²C");
        if (path.contains("/serio"))
            return QStringLiteral("Input");
        if (path.contains("/platform"))
            return QStringLiteral("Platform");
        return QStringLiteral("Other");
    }

    QList<WakeDevice> scanDevices()
    {
        QList<WakeDevice> result;
        QDirIterator it(QStringLiteral("/sys/devices"),
                        {QStringLiteral("wakeup")},
                        QDir::Files,
                        QDirIterator::Subdirectories);

        while (it.hasNext()) {
            const QString wakeFile = it.next();
            // wakeFile is .../<device>/power/wakeup. The device is two levels up
            // from the file, not the immediate parent named "power".
            QDir deviceDir(QFileInfo(wakeFile).absolutePath());
            if (!deviceDir.cdUp())
                continue;
            const QString path = deviceDir.absolutePath();
            const QString state = readText(wakeFile);
            if (state != "enabled" && state != "disabled")
                continue;

            WakeDevice device;
            device.wakeFile = wakeFile;
            device.sysfsPath = path;
            device.name = friendlyName(path);
            device.kind = deviceKind(path);
            device.enabled = state == "enabled";
            result.append(device);
        }

        std::sort(result.begin(), result.end(), [](const WakeDevice &a, const WakeDevice &b) {
            if (a.enabled != b.enabled)
                return a.enabled > b.enabled;
            return QString::localeAwareCompare(a.name, b.name) < 0;
        });
        return result;
    }

    class WakeControlWindow final : public QWidget
    {
    public:
        WakeControlWindow()
        {
            setWindowTitle(QStringLiteral("Wake Control"));
            resize(1050, 650);

            auto *root = new QVBoxLayout(this);
            auto *intro = new QLabel(
                QStringLiteral("Current wake permissions exposed by Linux. Changes are applied only when you choose an action."),
                                     this);
            intro->setWordWrap(true);
            root->addWidget(intro);

            auto *actions = new QHBoxLayout;
            refreshButton = new QPushButton(QStringLiteral("Refresh"), this);
            toggleAllButton = new QPushButton(QStringLiteral("Toggle all"), this);
            enableAllButton = new QPushButton(QStringLiteral("Enable all"), this);
            disableAllButton = new QPushButton(QStringLiteral("Disable all"), this);
            actions->addWidget(refreshButton);
            actions->addStretch();
            actions->addWidget(toggleAllButton);
            actions->addWidget(enableAllButton);
            actions->addWidget(disableAllButton);
            root->addLayout(actions);

            table = new QTableWidget(this);
            table->setColumnCount(5);
            table->setHorizontalHeaderLabels({QStringLiteral("Device"), QStringLiteral("Type"),
                QStringLiteral("Current state"), QStringLiteral("Pending"),
                                             QStringLiteral("Action")});
            table->setSelectionMode(QAbstractItemView::NoSelection);
            table->setEditTriggers(QAbstractItemView::NoEditTriggers);
            table->verticalHeader()->setVisible(false);
            table->horizontalHeader()->setStretchLastSection(true);
            table->horizontalHeader()->setSectionResizeMode(0, QHeaderView::Stretch);
            table->horizontalHeader()->setSectionResizeMode(1, QHeaderView::ResizeToContents);
            table->horizontalHeader()->setSectionResizeMode(2, QHeaderView::ResizeToContents);
            table->horizontalHeader()->setSectionResizeMode(3, QHeaderView::ResizeToContents);
            table->horizontalHeader()->setSortIndicatorShown(true);
            root->addWidget(table, 1);

            status = new QLabel(this);
            status->setStyleSheet(QStringLiteral("color: palette(mid);"));
            root->addWidget(status);

            auto *bottomActions = new QHBoxLayout;
            bottomActions->addStretch();
            cancelButton = new QPushButton(QStringLiteral("Cancel"), this);
            applyButton = new QPushButton(QStringLiteral("Apply"), this);
            savePermanentlyButton = new QPushButton(QStringLiteral("Save Permanently"), this);
            bottomActions->addWidget(cancelButton);
            bottomActions->addWidget(applyButton);
            bottomActions->addWidget(savePermanentlyButton);
            root->addLayout(bottomActions);

            connect(refreshButton, &QPushButton::clicked, this, [this] { reload(true); });
            connect(toggleAllButton, &QPushButton::clicked, this, [this] {
                const bool allEnabled = !devices.isEmpty() && std::all_of(
                    devices.cbegin(), devices.cend(), [](const WakeDevice &d) { return d.pending; });
                stageAll(!allEnabled);
            });
            connect(enableAllButton, &QPushButton::clicked, this, [this] { stageAll(true); });
            connect(disableAllButton, &QPushButton::clicked, this, [this] { stageAll(false); });
            connect(cancelButton, &QPushButton::clicked, this, [this] { reload(false); });
            connect(applyButton, &QPushButton::clicked, this, [this] { apply(); });
            connect(savePermanentlyButton, &QPushButton::clicked, this, [this] { savePermanently(); });
            connect(table->horizontalHeader(), &QHeaderView::sectionClicked, this, [this](int column) {
                if (sortColumn == column)
                    sortAscending = !sortAscending;
                else {
                    sortColumn = column;
                    sortAscending = true;
                }
                sortDevices();
                render();
            });

            reload(false);
        }

    private:
        bool hasPendingChanges() const
        {
            return std::any_of(devices.cbegin(), devices.cend(), [](const WakeDevice &device) {
                return device.enabled != device.pending;
            });
        }

        void reload(bool askBeforeDiscarding)
        {
            if (askBeforeDiscarding && hasPendingChanges()) {
                const auto answer = QMessageBox::question(
                    this, QStringLiteral("Discard unsaved changes?"),
                                                          QStringLiteral("Refresh will discard the wake-permission changes that have not been saved."),
                                                          QMessageBox::Discard | QMessageBox::Cancel, QMessageBox::Cancel);
                if (answer != QMessageBox::Discard)
                    return;
            }

            devices = scanDevices();
            for (WakeDevice &device : devices)
                device.pending = device.enabled;
            sortDevices();
            render();
        }

        void sortDevices()
        {
            std::stable_sort(devices.begin(), devices.end(), [this](const WakeDevice &left,
                                                                    const WakeDevice &right) {
                const auto value = [this](const WakeDevice &device) {
                    switch (sortColumn) {
                        case 0:
                            return device.name;
                        case 1:
                            return device.kind;
                        case 2:
                            return device.enabled ? QStringLiteral("Enabled") : QStringLiteral("Disabled");
                        case 3:
                            return device.pending ? QStringLiteral("Enabled") : QStringLiteral("Disabled");
                        case 4:
                            return device.pending ? QStringLiteral("Disable") : QStringLiteral("Enable");
                        default:
                            return QString();
                    }
                };

                const int comparison = QString::localeAwareCompare(value(left), value(right));
                if (comparison == 0)
                    return false;
                return sortAscending ? comparison < 0 : comparison > 0;
                                                                    });
        }

        void render()
        {
            table->horizontalHeader()->setSortIndicator(
                sortColumn, sortAscending ? Qt::AscendingOrder : Qt::DescendingOrder);
            table->setRowCount(0);
            table->setRowCount(devices.size());

            for (int row = 0; row < devices.size(); ++row) {
                const WakeDevice &device = devices.at(row);
                auto *name = new QTableWidgetItem(device.name);
                name->setToolTip(device.sysfsPath);
                table->setItem(row, 0, name);
                table->setItem(row, 1, new QTableWidgetItem(device.kind));

                auto *state = new QTableWidgetItem(device.enabled ? QStringLiteral("Enabled")
                : QStringLiteral("Disabled"));
                state->setTextAlignment(Qt::AlignCenter);
                table->setItem(row, 2, state);

                auto *pending = new QTableWidgetItem(device.pending ? QStringLiteral("Enabled")
                : QStringLiteral("Disabled"));
                pending->setTextAlignment(Qt::AlignCenter);
                if (device.pending != device.enabled)
                    pending->setText(pending->text() + QStringLiteral(" *"));
                table->setItem(row, 3, pending);

                auto *button = new QPushButton(device.pending ? QStringLiteral("Disable")
                : QStringLiteral("Enable"), table);
                connect(button, &QPushButton::clicked, this, [this, row] {
                    if (row >= 0 && row < devices.size())
                        stageOne(row, !devices.at(row).pending);
                });
                table->setCellWidget(row, 4, button);
            }

            const int enabled = std::count_if(devices.cbegin(), devices.cend(), [](const WakeDevice &d) {
                return d.enabled;
            });
            const int pendingChanges = std::count_if(devices.cbegin(), devices.cend(), [](const WakeDevice &d) {
                return d.enabled != d.pending;
            });
            status->setText(QStringLiteral("%1 devices found · %2 currently enabled · %3 unsaved change(s)")
            .arg(devices.size()).arg(enabled).arg(pendingChanges));
            applyButton->setEnabled(pendingChanges > 0);
            cancelButton->setEnabled(pendingChanges > 0);
            savePermanentlyButton->setEnabled(!devices.isEmpty());
        }

        bool runPrivileged(const QString &program, const QStringList &arguments,
                           const QByteArray &input = {}, QString *error = nullptr)
        {
            QProcess process;
            process.setProgram(QStringLiteral("pkexec"));
            QStringList pkexecArguments;
            pkexecArguments.append(program);
            pkexecArguments.append(arguments);
            process.setArguments(pkexecArguments);
            process.start();
            if (!process.waitForStarted(3000)) {
                if (error)
                    *error = process.errorString();
                return false;
            }
            if (!input.isEmpty())
                process.write(input);
            process.closeWriteChannel();
            if (!process.waitForFinished(15000)) {
                process.kill();
                if (error)
                    *error = QStringLiteral("Timed out");
                return false;
            }
            const bool succeeded = process.exitStatus() == QProcess::NormalExit && process.exitCode() == 0;
            if (!succeeded && error) {
                *error = QString::fromLocal8Bit(process.readAllStandardError()).trimmed();
                if (error->isEmpty())
                    *error = QStringLiteral("The administrator action was cancelled or failed.");
            }
            return succeeded;
        }

        bool writeState(const QString &path, bool enabled)
        {
            // pkexec invokes the desktop polkit agent, which provides the administrator prompt.
            return runPrivileged(QStringLiteral("/usr/bin/tee"), {path},
                                 enabled ? QByteArrayLiteral("enabled\n")
                                 : QByteArrayLiteral("disabled\n"));
        }

        void stageOne(int row, bool enabled)
        {
            devices[row].pending = enabled;
            render();
        }

        void stageAll(bool enabled)
        {
            for (WakeDevice &device : devices)
                device.pending = enabled;
            render();
        }

        QString udevQuoted(const QString &value) const
        {
            QString escaped = value;
            escaped.replace(QStringLiteral("\\"), QStringLiteral("\\\\"));
            escaped.replace(QStringLiteral("\""), QStringLiteral("\\\""));
            return escaped;
        }

        QString persistentRules() const
        {
            QString rules;
            rules += QStringLiteral("# Generated by Wake Control.\n");
            rules += QStringLiteral("# These rules preserve the current choices for the devices shown by the app.\n");
            rules += QStringLiteral("# They do not impose a default policy on unrelated future devices.\n\n");

            for (const WakeDevice &device : devices) {
                QString devPath = device.sysfsPath;
                if (devPath.startsWith(QStringLiteral("/sys")))
                    devPath = devPath.mid(4);
                if (devPath.isEmpty() || !devPath.startsWith(QLatin1Char('/')))
                    continue;

                rules += QStringLiteral("ACTION==\"add|change|bind\", DEVPATH==\"%1\", "
                "TEST==\"power/wakeup\", ATTR{power/wakeup}=\"%2\"\n")
                .arg(udevQuoted(devPath), device.pending ? QStringLiteral("enabled")
                : QStringLiteral("disabled"));
            }
            return rules;
        }

        bool applyPendingStates(QStringList *failed)
        {
            for (const WakeDevice &device : std::as_const(devices)) {
                if (device.enabled == device.pending)
                    continue;
                if (!writeState(device.wakeFile, device.pending) && failed)
                    failed->append(device.sysfsPath);
            }
            return !failed || failed->isEmpty();
        }

        void apply()
        {
            if (!hasPendingChanges())
                return;

            const int changes = std::count_if(devices.cbegin(), devices.cend(), [](const WakeDevice &device) {
                return device.enabled != device.pending;
            });
            const auto answer = QMessageBox::question(
                this, QStringLiteral("Apply wake-permission changes?"),
                                                      QStringLiteral("Apply %1 change(s) to the running kernel?\n\n"
                                                      "These changes are temporary and may be reset by a reboot or device reconnect.")
                                                      .arg(changes),
                                                      QMessageBox::Apply | QMessageBox::Cancel, QMessageBox::Apply);
            if (answer != QMessageBox::Apply)
                return;

            QStringList failed;
            applyPendingStates(&failed);
            reload(false);
            if (!failed.isEmpty()) {
                QMessageBox::warning(this, QStringLiteral("Some changes could not be applied"),
                                     QStringLiteral("%1 device(s) could not be changed.\n\n%2")
                                     .arg(failed.size())
                                     .arg(failed.join(QLatin1Char('\n'))));
            }
        }

        void savePermanently()
        {
            const int pendingChanges = std::count_if(devices.cbegin(), devices.cend(), [](const WakeDevice &device) {
                return device.enabled != device.pending;
            });
            const auto answer = QMessageBox::question(
                this, QStringLiteral("Save wake permissions permanently?"),
                                                      QStringLiteral("Persist the current pending state for all %1 devices?\n\n"
                                                      "This will open an administrator authentication window and write a udev rule. "
                                                      "The settings will also be applied to the running kernel%2.")
                                                      .arg(devices.size())
                                                      .arg(pendingChanges > 0 ? QStringLiteral(".") : QStringLiteral(" (there are no pending changes)")),
                                                      QMessageBox::Save | QMessageBox::Cancel, QMessageBox::Save);
            if (answer != QMessageBox::Save)
                return;

            QTemporaryFile temporaryFile(QDir::tempPath() + QStringLiteral("/wake-control-XXXXXX.rules"));
            temporaryFile.setAutoRemove(false);
            if (!temporaryFile.open()) {
                QMessageBox::critical(this, QStringLiteral("Cannot save permanent policy"),
                                      QStringLiteral("Could not create a temporary rules file."));
                return;
            }

            const QByteArray rules = persistentRules().toUtf8();
            if (temporaryFile.write(rules) != rules.size() || !temporaryFile.flush()) {
                const QString path = temporaryFile.fileName();
                temporaryFile.remove();
                QMessageBox::critical(this, QStringLiteral("Cannot save permanent policy"),
                                      QStringLiteral("Could not write the temporary rules file:\n%1").arg(path));
                return;
            }
            const QString temporaryPath = temporaryFile.fileName();
            temporaryFile.close();

            QString error;
            const bool installed = runPrivileged(
                QStringLiteral("/usr/bin/install"),
                                                 {QStringLiteral("-o"), QStringLiteral("root"), QStringLiteral("-g"), QStringLiteral("root"),
                                                     QStringLiteral("-m"), QStringLiteral("0644"), temporaryPath,
                                                 QStringLiteral("/etc/udev/rules.d/99-wakeup-control.rules")}, {}, &error);
            QFile::remove(temporaryPath);

            if (!installed) {
                QMessageBox::warning(this, QStringLiteral("Permanent save cancelled"),
                                     QStringLiteral("The administrator action did not install the udev rule.\n\n%1")
                                     .arg(error));
                return;
            }

            // Reload rules so the next device event uses the new policy. Current devices
            // are applied explicitly below because udev rules do not retroactively run.
            QString reloadError;
            const bool reloaded = runPrivileged(QStringLiteral("/usr/bin/udevadm"), {QStringLiteral("control"),
                QStringLiteral("--reload-rules")},
                                                {}, &reloadError);

            QStringList failed;
            applyPendingStates(&failed);
            reload(false);

            if (!failed.isEmpty()) {
                QMessageBox::warning(this, QStringLiteral("Permanent policy saved; some runtime changes failed"),
                                     QStringLiteral("The persistent udev rule was installed, but %1 current device(s) "
                                     "could not be changed:\n\n%2")
                                     .arg(failed.size())
                                     .arg(failed.join(QLatin1Char('\n'))));
            } else if (!reloaded) {
                QMessageBox::warning(this, QStringLiteral("Permanent policy saved; udev reload failed"),
                                     QStringLiteral("The persistent rule was installed and current settings were applied, "
                                     "but udev could not be reloaded immediately. It will take effect "
                                     "after the next udev reload or reboot.\n\n%1")
                                     .arg(reloadError));
            } else {
                QMessageBox::information(this, QStringLiteral("Wake policy saved"),
                                         QStringLiteral("The wake-permission policy was saved permanently and applied now."));
            }
        }

        QList<WakeDevice> devices;
        QTableWidget *table = nullptr;
        QLabel *status = nullptr;
        QPushButton *refreshButton = nullptr;
        QPushButton *toggleAllButton = nullptr;
        QPushButton *enableAllButton = nullptr;
        QPushButton *disableAllButton = nullptr;
        QPushButton *cancelButton = nullptr;
        QPushButton *applyButton = nullptr;
        QPushButton *savePermanentlyButton = nullptr;
        int sortColumn = 0;
        bool sortAscending = true;
    };

} // namespace

int main(int argc, char **argv)
{
    QApplication app(argc, argv);
    QApplication::setApplicationName(QStringLiteral("Wake Control"));
    QApplication::setApplicationVersion(QStringLiteral("0.1.0"));
    WakeControlWindow window;
    window.show();
    return app.exec();
}