WIP: implement avahi service

src/jailbrokenwidget.cpp

@@ -1,15 +1,21 @@
 #include "jailbrokenwidget.h"
 #include "appcontext.h"
+#include "core/services/avahi_service.h"
 #include "iDescriptor-ui.h"
 #include "iDescriptor.h"
+#include <QButtonGroup>
 #include <QDebug>
 #include <QGraphicsPixmapItem>
 #include <QGraphicsView>
+#include <QGroupBox>
 #include <QHBoxLayout>
 #include <QLabel>
+#include <QMenu>
 #include <QProcess>
 #include <QProcessEnvironment>
 #include <QPushButton>
+#include <QRadioButton>
+#include <QScrollArea>
 #include <QTimer>
 #include <QVBoxLayout>
 #include <QWidget>
@@ -36,10 +42,20 @@ JailbrokenWidget::JailbrokenWidget(QWidget *parent) : QWidget{parent}
 
     mainLayout->addWidget(graphicsView, 1);
 
+    // Connect to AppContext for device events
     connect(AppContext::sharedInstance(), &AppContext::deviceAdded, this,
-            [this](iDescriptorDevice *device) { deviceConnected(device); });
+            &JailbrokenWidget::onWiredDeviceAdded);
+    connect(AppContext::sharedInstance(), &AppContext::deviceRemoved, this,
+            &JailbrokenWidget::onWiredDeviceRemoved);
 
-    // Right side: Info and Terminal
+    // Initialize Avahi service
+    m_avahiService = new AvahiService(this);
+    connect(m_avahiService, &AvahiService::deviceAdded, this,
+            &JailbrokenWidget::onWirelessDeviceAdded);
+    connect(m_avahiService, &AvahiService::deviceRemoved, this,
+            &JailbrokenWidget::onWirelessDeviceRemoved);
+
+    // Right side: Device selection and Terminal
     QWidget *rightContainer = new QWidget();
     rightContainer->setSizePolicy(QSizePolicy::Expanding,
                                   QSizePolicy::Expanding);
@@ -48,15 +64,7 @@ JailbrokenWidget::JailbrokenWidget(QWidget *parent) : QWidget{parent}
     rightLayout->setContentsMargins(15, 15, 15, 15);
     rightLayout->setSpacing(10);
 
-    m_infoLabel = new QLabel("Connect a jailbroken device");
-    rightLayout->addWidget(m_infoLabel);
-
-    m_connectButton = new QPushButton("Connect SSH Terminal");
-    m_connectButton->setEnabled(false);
-    connect(m_connectButton, &QPushButton::clicked, this,
-            &JailbrokenWidget::onConnectSSH);
-    rightLayout->addWidget(m_connectButton);
-
+    setupDeviceSelectionUI(rightLayout);
     setupTerminal();
     rightLayout->addWidget(m_terminal, 1);
 
@@ -71,6 +79,12 @@ JailbrokenWidget::JailbrokenWidget(QWidget *parent) : QWidget{parent}
     m_sshTimer = new QTimer(this);
     connect(m_sshTimer, &QTimer::timeout, this,
             &JailbrokenWidget::checkSshData);
+
+    // Start scanning for wireless devices
+    m_avahiService->startBrowsing();
+
+    // Populate initial devices
+    updateDeviceList();
 }
 
 void JailbrokenWidget::setupTerminal()
@@ -78,9 +92,228 @@ void JailbrokenWidget::setupTerminal()
     m_terminal = new QTermWidget(0, this);
     m_terminal->setMinimumHeight(400);
     m_terminal->setScrollBarPosition(QTermWidget::ScrollBarRight);
-    m_terminal->setColorScheme("DarkPastels");
+    m_terminal->setColorScheme("Linux");
+    m_terminal->setContextMenuPolicy(Qt::CustomContextMenu);
+    connect(m_terminal, &QWidget::customContextMenuRequested, this,
+            [this](const QPoint &pos) {
+                QMenu menu(this);
+                QList<QAction *> actions = m_terminal->filterActions(pos);
+                if (!actions.isEmpty()) {
+                    menu.addActions(actions);
+                    menu.exec(m_terminal->mapToGlobal(pos));
+                }
+            });
     m_terminal->startTerminalTeletype();
     m_terminal->hide();
+    m_terminal->setStyleSheet("padding : 10px;");
+}
+
+void JailbrokenWidget::setupDeviceSelectionUI(QVBoxLayout *layout)
+{
+    // Create scroll area for device selection
+    QScrollArea *scrollArea = new QScrollArea();
+    scrollArea->setWidgetResizable(true);
+    scrollArea->setMinimumHeight(200);
+    scrollArea->setMaximumHeight(300);
+
+    QWidget *scrollContent = new QWidget();
+    m_deviceLayout = new QVBoxLayout(scrollContent);
+    m_deviceLayout->setContentsMargins(5, 5, 5, 5);
+    m_deviceLayout->setSpacing(10);
+
+    // Button group for device selection
+    m_deviceButtonGroup = new QButtonGroup(this);
+    connect(m_deviceButtonGroup,
+            QOverload<QAbstractButton *>::of(&QButtonGroup::buttonClicked),
+            this, &JailbrokenWidget::onDeviceSelected);
+
+    // Wired devices group
+    m_wiredDevicesGroup = new QGroupBox("Connected Devices");
+    m_wiredDevicesLayout = new QVBoxLayout(m_wiredDevicesGroup);
+    m_deviceLayout->addWidget(m_wiredDevicesGroup);
+
+    // Wireless devices group
+    m_wirelessDevicesGroup = new QGroupBox("Network Devices");
+    m_wirelessDevicesLayout = new QVBoxLayout(m_wirelessDevicesGroup);
+    m_deviceLayout->addWidget(m_wirelessDevicesGroup);
+
+    scrollArea->setWidget(scrollContent);
+    layout->addWidget(scrollArea);
+
+    // Info and connect button
+    m_infoLabel = new QLabel("Select a device to connect");
+    layout->addWidget(m_infoLabel);
+
+    m_connectButton = new QPushButton("Connect SSH Terminal");
+    m_connectButton->setEnabled(false);
+    connect(m_connectButton, &QPushButton::clicked, this,
+            &JailbrokenWidget::onConnectSSH);
+    layout->addWidget(m_connectButton);
+}
+
+void JailbrokenWidget::updateDeviceList()
+{
+    // Clear existing devices
+    clearDeviceButtons();
+
+    // Add wired devices
+    QList<iDescriptorDevice *> wiredDevices =
+        AppContext::sharedInstance()->getAllDevices();
+    for (iDescriptorDevice *device : wiredDevices) {
+        addWiredDevice(device);
+    }
+
+    // Add wireless devices
+    QList<NetworkDevice> wirelessDevices = m_avahiService->getNetworkDevices();
+    for (const NetworkDevice &device : wirelessDevices) {
+        addWirelessDevice(device);
+    }
+}
+
+void JailbrokenWidget::clearDeviceButtons()
+{
+    // Remove all buttons from button group and layouts
+    for (QAbstractButton *button : m_deviceButtonGroup->buttons()) {
+        m_deviceButtonGroup->removeButton(button);
+        button->deleteLater();
+    }
+
+    // Clear layouts
+    QLayoutItem *item;
+    while ((item = m_wiredDevicesLayout->takeAt(0)) != nullptr) {
+        delete item->widget();
+        delete item;
+    }
+    while ((item = m_wirelessDevicesLayout->takeAt(0)) != nullptr) {
+        delete item->widget();
+        delete item;
+    }
+}
+
+void JailbrokenWidget::addWiredDevice(iDescriptorDevice *device)
+{
+    QString deviceName = QString::fromStdString(device->deviceInfo.deviceName);
+    QString udid = QString::fromStdString(device->udid);
+    QString displayText = QString("%1\n%2").arg(deviceName, udid);
+
+    QRadioButton *radioButton = new QRadioButton(displayText);
+    radioButton->setProperty("deviceType", "wired");
+    radioButton->setProperty("devicePointer",
+                             QVariant::fromValue(static_cast<void *>(device)));
+    radioButton->setProperty("udid", udid);
+
+    m_deviceButtonGroup->addButton(radioButton);
+    m_wiredDevicesLayout->addWidget(radioButton);
+}
+
+void JailbrokenWidget::addWirelessDevice(const NetworkDevice &device)
+{
+    QString displayText = QString("%1\n%2").arg(device.name, device.address);
+
+    QRadioButton *radioButton = new QRadioButton(displayText);
+    radioButton->setProperty("deviceType", "wireless");
+    radioButton->setProperty("deviceAddress", device.address);
+    radioButton->setProperty("deviceName", device.name);
+    radioButton->setProperty("devicePort", device.port);
+
+    m_deviceButtonGroup->addButton(radioButton);
+    m_wirelessDevicesLayout->addWidget(radioButton);
+}
+
+void JailbrokenWidget::onWiredDeviceAdded(iDescriptorDevice *device)
+{
+    addWiredDevice(device);
+}
+
+void JailbrokenWidget::onWiredDeviceRemoved(const std::string &udid)
+{
+    QString qudid = QString::fromStdString(udid);
+
+    // Find and remove the corresponding radio button
+    for (QAbstractButton *button : m_deviceButtonGroup->buttons()) {
+        if (button->property("deviceType").toString() == "wired" &&
+            button->property("udid").toString() == qudid) {
+            m_deviceButtonGroup->removeButton(button);
+            button->deleteLater();
+            break;
+        }
+    }
+
+    // Reset selection if this device was selected
+    if (m_selectedDeviceType == DeviceType::Wired && m_selectedWiredDevice &&
+        m_selectedWiredDevice->udid == udid) {
+        resetSelection();
+    }
+}
+
+void JailbrokenWidget::onWirelessDeviceAdded(const NetworkDevice &device)
+{
+    addWirelessDevice(device);
+}
+
+void JailbrokenWidget::onWirelessDeviceRemoved(const QString &deviceName)
+{
+    // Find and remove the corresponding radio button
+    for (QAbstractButton *button : m_deviceButtonGroup->buttons()) {
+        if (button->property("deviceType").toString() == "wireless" &&
+            button->property("deviceName").toString() == deviceName) {
+            m_deviceButtonGroup->removeButton(button);
+            button->deleteLater();
+            break;
+        }
+    }
+
+    // Reset selection if this device was selected
+    if (m_selectedDeviceType == DeviceType::Wireless &&
+        m_selectedNetworkDevice.name == deviceName) {
+        resetSelection();
+    }
+}
+
+void JailbrokenWidget::onDeviceSelected(QAbstractButton *button)
+{
+    QString deviceType = button->property("deviceType").toString();
+
+    if (deviceType == "wired") {
+        m_selectedDeviceType = DeviceType::Wired;
+        m_selectedWiredDevice = static_cast<iDescriptorDevice *>(
+            button->property("devicePointer").value<void *>());
+
+        if (m_selectedWiredDevice->deviceInfo.jailbroken) {
+            m_infoLabel->setText("Jailbroken device selected");
+        } else {
+            m_infoLabel->setText("Device selected (jailbreak status unknown)");
+        }
+    } else if (deviceType == "wireless") {
+        m_selectedDeviceType = DeviceType::Wireless;
+        m_selectedNetworkDevice.name =
+            button->property("deviceName").toString();
+        m_selectedNetworkDevice.address =
+            button->property("deviceAddress").toString();
+        m_selectedNetworkDevice.port = button->property("devicePort").toUInt();
+
+        m_infoLabel->setText(
+            "Network device selected (jailbreak status unknown)");
+    }
+
+    m_connectButton->setEnabled(true);
+    m_connectButton->setText("Connect SSH Terminal");
+}
+
+void JailbrokenWidget::resetSelection()
+{
+    m_selectedDeviceType = DeviceType::None;
+    m_selectedWiredDevice = nullptr;
+    m_selectedNetworkDevice = NetworkDevice{};
+    m_connectButton->setEnabled(false);
+    m_infoLabel->setText("Select a device to connect");
+
+    // Uncheck all radio buttons
+    if (m_deviceButtonGroup->checkedButton()) {
+        m_deviceButtonGroup->setExclusive(false);
+        m_deviceButtonGroup->checkedButton()->setChecked(false);
+        m_deviceButtonGroup->setExclusive(true);
+    }
 }
 
 void JailbrokenWidget::connectLibsshToTerminal()
@@ -100,11 +333,14 @@ void JailbrokenWidget::connectLibsshToTerminal()
 void JailbrokenWidget::deviceConnected(iDescriptorDevice *device)
 {
     if (device->deviceInfo.jailbroken) {
-        m_device = device;
         m_infoLabel->setText("Jailbroken device connected");
-        m_connectButton->setEnabled(true);
-        m_connectButton->setText("Connect SSH Terminal");
+    } else {
+        m_infoLabel->setText(
+            "Connected device is not detected as jailbroken. Continue anyway?");
     }
+    m_device = device;
+    m_connectButton->setEnabled(true);
+    m_connectButton->setText("Connect SSH Terminal");
 }
 
 void JailbrokenWidget::onConnectSSH()
@@ -114,24 +350,33 @@ void JailbrokenWidget::onConnectSSH()
         return;
     }
 
-    initWidget();
+    if (m_selectedDeviceType == DeviceType::None) {
+        m_infoLabel->setText("Please select a device first");
+        return;
+    }
+
+    if (m_selectedDeviceType == DeviceType::Wired) {
+        initWiredDevice();
+    } else {
+        initWirelessDevice();
+    }
 }
 
-void JailbrokenWidget::initWidget()
+void JailbrokenWidget::initWiredDevice()
 {
     if (m_isInitialized)
         return;
     m_isInitialized = true;
 
-    if (!m_device) {
-        m_infoLabel->setText("Device is not jailbroken");
+    if (!m_selectedWiredDevice) {
+        m_infoLabel->setText("No wired device selected");
         return;
     }
 
     m_connectButton->setEnabled(false);
     m_infoLabel->setText("Setting up SSH tunnel...");
 
-    // Start iproxy first
+    // Start iproxy first for wired devices
     iproxyProcess = new QProcess(this);
     iproxyProcess->setProcessChannelMode(QProcess::MergedChannels);
     QProcessEnvironment env = QProcessEnvironment::systemEnvironment();
@@ -147,13 +392,22 @@ void JailbrokenWidget::initWidget()
                 qDebug() << "iproxy error:" << error;
             });
 
-    QTimer::singleShot(
-        3000, this,
-        &JailbrokenWidget::startSSH); // Increased delay to 3 seconds
+    // Monitor iproxy output for readiness
+    connect(iproxyProcess, &QProcess::readyRead, this, [this]() {
+        QByteArray output = iproxyProcess->readAll();
+        if (output.contains("waiting for connection")) {
+            // iproxy is ready, disconnect the signal to avoid multiple calls
+            disconnect(iproxyProcess, &QProcess::readyRead, this, nullptr);
+            startSSH("127.0.0.1", 3333);
+        }
+    });
 
-    iproxyProcess->start("iproxy", QStringList()
-                                       << "-u" << m_device->udid.c_str()
-                                       << "3333" << "22");
+    QStringList args;
+    args << "-u" << m_selectedWiredDevice->udid.c_str() << "3333" << "22";
+
+    qDebug() << "Starting iproxy with args:" << args;
+
+    iproxyProcess->start("iproxy", args);
 
     // Check if iproxy started successfully
     if (!iproxyProcess->waitForStarted(5000)) {
@@ -163,13 +417,26 @@ void JailbrokenWidget::initWidget()
     }
 }
 
-void JailbrokenWidget::startSSH()
+void JailbrokenWidget::initWirelessDevice()
+{
+    if (m_isInitialized)
+        return;
+    m_isInitialized = true;
+
+    m_connectButton->setEnabled(false);
+    m_infoLabel->setText("Connecting to network device...");
+
+    // For wireless devices, connect directly without iproxy
+    startSSH(m_selectedNetworkDevice.address, m_selectedNetworkDevice.port);
+}
+
+void JailbrokenWidget::startSSH(const QString &host, uint16_t port)
 {
     if (m_sshConnected)
         return;
 
     m_infoLabel->setText("Connecting to SSH server...");
-    qDebug() << "Starting SSH connection to localhost:3333";
+    qDebug() << "Starting SSH connection to" << host << ":" << port;
 
     // Create SSH session
     m_sshSession = ssh_new();
@@ -180,9 +447,10 @@ void JailbrokenWidget::startSSH()
     }
 
     // Configure SSH session
-    ssh_options_set(m_sshSession, SSH_OPTIONS_HOST, "localhost");
-    int port = 3333;
-    ssh_options_set(m_sshSession, SSH_OPTIONS_PORT, &port);
+    QByteArray hostBytes = host.toUtf8();
+    ssh_options_set(m_sshSession, SSH_OPTIONS_HOST, hostBytes.constData());
+    int sshPort = static_cast<int>(port);
+    ssh_options_set(m_sshSession, SSH_OPTIONS_PORT, &sshPort);
     ssh_options_set(m_sshSession, SSH_OPTIONS_USER, "root");
 
     // Disable strict host key checking
@@ -345,8 +613,12 @@ void JailbrokenWidget::disconnectSSH()
     }
 
     if (iproxyProcess) {
-        iproxyProcess->terminate();
-        iproxyProcess->waitForFinished(3000);
+        iproxyProcess->terminate(); // Ask it to terminate nicely
+        if (!iproxyProcess->waitForFinished(1000)) { // Wait 1 sec
+            iproxyProcess->kill();                   // Forcefully kill it
+            iproxyProcess->waitForFinished(1000);    // Wait for it to die
+        }
+        delete iproxyProcess; // Avoid memory leak
         iproxyProcess = nullptr;
     }