[rtns] Add functionality to enable and disable webserver (#6270)

2026-06-19 00:41:43 +00:00 · 2024-01-20 20:15:30 -08:00
parent 7957f4a625
commit e408f3ad27
5 changed files with 336 additions and 32 deletions
--- a/roborioteamnumbersetter/src/main/native/cpp/App.cpp
+++ b/roborioteamnumbersetter/src/main/native/cpp/App.cpp
@@ -57,6 +57,8 @@ struct TeamNumberRefHolder {
 static std::unique_ptr<TeamNumberRefHolder> teamNumberRef;
 static std::unordered_map<std::string, std::pair<unsigned int, std::string>>
    foundDevices;
+static std::unordered_map<std::string, std::optional<sysid::DeviceStatus>>
+    deviceStatuses;
 static wpi::Logger logger;
 static sysid::DeploySession deploySession{logger};
 static std::unique_ptr<wpi::MulticastServiceResolver> multicastResolver;
@@ -76,7 +78,12 @@ static void FindDevices() {
                       [](const auto& a) { return a.first == "MAC"; });
      if (macKey != data.txt.end()) {
        auto& mac = macKey->second;
-        foundDevices[mac] = std::make_pair(data.ipv4Address, data.hostName);
+        auto& foundDevice = foundDevices[mac];
+        foundDevice = std::make_pair(data.ipv4Address, data.hostName);
+        auto& deviceStatus = deviceStatuses[mac];
+        if (!deviceStatus) {
+          deploySession.GetStatus(mac, foundDevice.first);
+        }
      }
    }
  }
@@ -146,15 +153,12 @@ static void DisplayGui() {
    int macWidth = ImGui::CalcTextSize("88:88:88:88:88:88").x;
    int ipAddressWidth = ImGui::CalcTextSize("255.255.255.255").x;
    int setWidth = ImGui::CalcTextSize(" Set Team To 99999 ").x;
-    int blinkWidth = ImGui::CalcTextSize(" Blink ").x;
-    int rebootWidth = ImGui::CalcTextSize(" Reboot ").x;

-    minWidth = nameWidth + macWidth + ipAddressWidth + setWidth + blinkWidth +
-               rebootWidth + 100;
+    minWidth = nameWidth + macWidth + ipAddressWidth + setWidth + 100;

    std::string setString = fmt::format("Set team to {}", teamNumber);

-    if (ImGui::BeginTable("Table", 6)) {
+    if (ImGui::BeginTable("Table", 4)) {
      ImGui::TableSetupColumn(
          "Name",
          ImGuiTableColumnFlags_NoResize | ImGuiTableColumnFlags_WidthFixed,
@@ -171,17 +175,33 @@ static void DisplayGui() {
          "Set",
          ImGuiTableColumnFlags_NoResize | ImGuiTableColumnFlags_WidthFixed,
          setWidth);
-      ImGui::TableSetupColumn(
-          "Blink",
-          ImGuiTableColumnFlags_NoResize | ImGuiTableColumnFlags_WidthFixed,
-          blinkWidth);
-      ImGui::TableSetupColumn(
-          "Reboot",
-          ImGuiTableColumnFlags_NoResize | ImGuiTableColumnFlags_WidthFixed,
-          rebootWidth);
      ImGui::TableHeadersRow();

-      for (auto&& i : foundDevices) {
+      ImGui::EndTable();
+    }
+
+    for (auto&& i : foundDevices) {
+      std::future<int>* future = deploySession.GetFuture(i.first);
+      std::future<sysid::DeviceStatus>* futureStatus =
+          deploySession.GetStatusFuture(i.first);
+      if (ImGui::BeginTable("Table", 4)) {
+        ImGui::TableSetupColumn(
+            "Name",
+            ImGuiTableColumnFlags_NoResize | ImGuiTableColumnFlags_WidthFixed,
+            nameWidth);
+        ImGui::TableSetupColumn(
+            "MAC Address",
+            ImGuiTableColumnFlags_NoResize | ImGuiTableColumnFlags_WidthFixed,
+            macWidth);
+        ImGui::TableSetupColumn(
+            "IP Address",
+            ImGuiTableColumnFlags_NoResize | ImGuiTableColumnFlags_WidthFixed,
+            ipAddressWidth);
+        ImGui::TableSetupColumn(
+            "Set",
+            ImGuiTableColumnFlags_NoResize | ImGuiTableColumnFlags_WidthFixed,
+            setWidth);
+
        ImGui::TableNextRow();
        ImGui::TableNextColumn();
        ImGui::Text("%s", i.second.second.c_str());
@@ -192,11 +212,8 @@ static void DisplayGui() {
        in.s_addr = i.second.first;
        ImGui::Text("%s", inet_ntoa(in));
        ImGui::TableNextColumn();
-        std::future<int>* future = deploySession.GetFuture(i.first);
        ImGui::PushID(i.first.c_str());
        if (future) {
-          ImGui::Button("Deploying");
-          ImGui::TableNextColumn();
          ImGui::TableNextColumn();
          const auto fs = future->wait_for(std::chrono::seconds(0));
          if (fs == std::future_status::ready) {
@@ -207,18 +224,62 @@ static void DisplayGui() {
            deploySession.ChangeTeamNumber(i.first, teamNumber, i.second.first);
          }
          ImGui::TableNextColumn();
-          if (ImGui::Button("Blink")) {
-            deploySession.Blink(i.first, i.second.first);
-          }
-          ImGui::TableNextColumn();
-          if (ImGui::Button("Reboot")) {
-            deploySession.Reboot(i.first, i.second.first);
-          }
        }
+
        ImGui::PopID();
+        ImGui::EndTable();
      }

-      ImGui::EndTable();
+      ImGui::PushID(i.first.c_str());
+      if (futureStatus) {
+        ImGui::Text("Refreshing Status");
+        const auto fs = futureStatus->wait_for(std::chrono::seconds(0));
+        if (fs == std::future_status::ready) {
+          deviceStatuses[i.first] = futureStatus->get();
+          deploySession.DestroyStatusFuture(i.first);
+        }
+      } else {
+        auto& deviceStatus = deviceStatuses[i.first];
+        if (deviceStatus) {
+          if (ImGui::Button("Refresh Status")) {
+            deploySession.GetStatus(i.first, i.second.first);
+          }
+          std::string formatted =
+              fmt::format("Image: {}", deviceStatus.value().image);
+          ImGui::Text("%s", formatted.c_str());
+          formatted = fmt::format("Serial Number: {}",
+                                  deviceStatus.value().serialNumber);
+          ImGui::Text("%s", formatted.c_str());
+          formatted = fmt::format(
+              "Web Server Status: {}",
+              deviceStatus.value().webServerEnabled ? "Enabled" : "Disabled");
+          ImGui::Text("%s", formatted.c_str());
+        } else {
+          ImGui::Text("Waiting for refresh");
+        }
+      }
+
+      if (future) {
+        ImGui::Text("Deploying");
+      } else {
+        if (ImGui::Button("Blink")) {
+          deploySession.Blink(i.first, i.second.first);
+        }
+        ImGui::SameLine();
+        if (ImGui::Button("Reboot")) {
+          deploySession.Reboot(i.first, i.second.first);
+        }
+        ImGui::SameLine();
+        if (ImGui::Button("Disable Web Server")) {
+          deploySession.DisableWebServer(i.first, i.second.first);
+        }
+        ImGui::SameLine();
+        if (ImGui::Button("Enable Web Server")) {
+          deploySession.EnableWebServer(i.first, i.second.first);
+        }
+      }
+      ImGui::Separator();
+      ImGui::PopID();
    }

    ImGui::Columns(6, "Devices");
--- a/roborioteamnumbersetter/src/main/native/cpp/DeploySession.cpp
+++ b/roborioteamnumbersetter/src/main/native/cpp/DeploySession.cpp
@@ -37,6 +37,7 @@ static constexpr std::string_view kUsername = "admin";
 static constexpr std::string_view kPassword = "";

 std::unordered_map<std::string, std::future<int>> s_outstanding;
+std::unordered_map<std::string, std::future<DeviceStatus>> s_outstandingStatus;

 DeploySession::DeploySession(wpi::Logger& logger) : m_logger{logger} {}

@@ -59,6 +60,19 @@ void DeploySession::DestroyFuture(const std::string& macAddress) {
  s_outstanding.erase(macAddress);
 }

+std::future<DeviceStatus>* DeploySession::GetStatusFuture(
+    const std::string& macAddress) {
+  auto itr = s_outstandingStatus.find(macAddress);
+  if (itr == s_outstandingStatus.end()) {
+    return nullptr;
+  }
+  return &itr->second;
+}
+
+void DeploySession::DestroyStatusFuture(const std::string& macAddress) {
+  s_outstandingStatus.erase(macAddress);
+}
+
 bool DeploySession::ChangeTeamNumber(const std::string& macAddress,
                                     int teamNumber, unsigned int ipAddress) {
  auto itr = s_outstanding.find(macAddress);
@@ -184,3 +198,147 @@ bool DeploySession::Blink(const std::string& macAddress,
  s_outstanding[macAddress] = std::move(future);
  return true;
 }
+
+bool DeploySession::DisableWebServer(const std::string& macAddress,
+                                     unsigned int ipAddress) {
+  auto itr = s_outstanding.find(macAddress);
+  if (itr != s_outstanding.end()) {
+    return false;
+  }
+
+  std::future<int> future =
+      std::async(std::launch::async, [this, ipAddress, mac = macAddress]() {
+        //  Convert to IP address.
+        wpi::SmallString<16> ip;
+        in_addr addr;
+        addr.s_addr = ipAddress;
+        wpi::uv::AddrToName(addr, &ip);
+        DEBUG("Trying to establish SSH connection to {}.", ip.str());
+        try {
+          SshSession session{ip.str(), kPort, kUsername, kPassword, m_logger};
+          session.Open();
+          DEBUG("SSH connection to {} was successful.", ip.str());
+
+          SUCCESS("roboRIO Connected!");
+
+          try {
+            session.Execute(
+                "/bin/bash -c \"/etc/init.d/systemWebServer stop\"");
+            session.Execute(
+                "/bin/bash -c \"update-rc.d -f systemWebServer remove\"");
+            session.Execute("/bin/bash -c \"sync\"");
+          } catch (const SshSession::SshException& e) {
+            ERROR("An exception occurred: {}", e.what());
+            throw e;
+          }
+        } catch (const SshSession::SshException& e) {
+          DEBUG("SSH connection to {} failed with {}.", ip.str(), e.what());
+          throw e;
+        }
+        return 0;
+      });
+
+  s_outstanding[macAddress] = std::move(future);
+  return true;
+}
+
+bool DeploySession::EnableWebServer(const std::string& macAddress,
+                                    unsigned int ipAddress) {
+  auto itr = s_outstanding.find(macAddress);
+  if (itr != s_outstanding.end()) {
+    return false;
+  }
+
+  std::future<int> future =
+      std::async(std::launch::async, [this, ipAddress, mac = macAddress]() {
+        //  Convert to IP address.
+        wpi::SmallString<16> ip;
+        in_addr addr;
+        addr.s_addr = ipAddress;
+        wpi::uv::AddrToName(addr, &ip);
+        DEBUG("Trying to establish SSH connection to {}.", ip.str());
+        try {
+          SshSession session{ip.str(), kPort, kUsername, kPassword, m_logger};
+          session.Open();
+          DEBUG("SSH connection to {} was successful.", ip.str());
+
+          SUCCESS("roboRIO Connected!");
+
+          try {
+            session.Execute(
+                "/bin/bash -c \"update-rc.d -f systemWebServer defaults\"");
+            session.Execute(
+                "/bin/bash -c \"/etc/init.d/systemWebServer start\"");
+            session.Execute("/bin/bash -c \"sync\"");
+          } catch (const SshSession::SshException& e) {
+            ERROR("An exception occurred: {}", e.what());
+            throw e;
+          }
+        } catch (const SshSession::SshException& e) {
+          DEBUG("SSH connection to {} failed with {}.", ip.str(), e.what());
+          throw e;
+        }
+        return 0;
+      });
+
+  s_outstanding[macAddress] = std::move(future);
+  return true;
+}
+
+bool DeploySession::GetStatus(const std::string& macAddress,
+                              unsigned int ipAddress) {
+  auto itr = s_outstandingStatus.find(macAddress);
+  if (itr != s_outstandingStatus.end()) {
+    return false;
+  }
+
+  std::future<DeviceStatus> future =
+      std::async(std::launch::async, [this, ipAddress, mac = macAddress]() {
+        //  Convert to IP address.
+        wpi::SmallString<16> ip;
+        in_addr addr;
+        addr.s_addr = ipAddress;
+        wpi::uv::AddrToName(addr, &ip);
+        DEBUG("Trying to establish SSH connection to {}.", ip.str());
+        DeviceStatus status;
+        try {
+          SshSession session{ip.str(), kPort, kUsername, kPassword, m_logger};
+          session.Open();
+          DEBUG("SSH connection to {} was successful.", ip.str());
+
+          SUCCESS("roboRIO Connected!");
+
+          try {
+            int exitStatus = 0;
+            session.ExecuteResult(
+                "start-stop-daemon --status -x "
+                "/usr/local/natinst/share/NIWebServer/SystemWebServer",
+                &exitStatus);
+            status.webServerEnabled = exitStatus == 0;
+            auto serialNumber = session.ExecuteResult(
+                "/sbin/fw_printenv -n serial#", &exitStatus);
+            if (exitStatus == 0) {
+              status.serialNumber = wpi::trim(serialNumber);
+            }
+            auto image = session.ExecuteResult(
+                "/usr/local/natinst/bin/nirtcfg --file "
+                "/etc/natinst/share/scs_imagemetadata.ini --get "
+                "section=ImageMetadata,token=IMAGEVERSION,value=UNKNOWN",
+                &exitStatus);
+            if (exitStatus == 0) {
+              status.image = wpi::trim(image);
+            }
+          } catch (const SshSession::SshException& e) {
+            ERROR("An exception occurred: {}", e.what());
+            throw e;
+          }
+        } catch (const SshSession::SshException& e) {
+          DEBUG("SSH connection to {} failed with {}.", ip.str(), e.what());
+          throw e;
+        }
+        return status;
+      });
+
+  s_outstandingStatus[macAddress] = std::move(future);
+  return true;
+}
--- a/roborioteamnumbersetter/src/main/native/cpp/DeploySession.h
+++ b/roborioteamnumbersetter/src/main/native/cpp/DeploySession.h
@@ -18,6 +18,12 @@ namespace sysid {
 // GUI).
 static constexpr unsigned int kLogSuccess = 31;

+struct DeviceStatus {
+  bool webServerEnabled = false;
+  std::string serialNumber;
+  std::string image;
+};
+
 /**
 * Represents a single deploy session.
 *
@@ -49,15 +55,19 @@ class DeploySession {

  bool Blink(const std::string& macAddress, unsigned int ipAddress);

+  bool DisableWebServer(const std::string& macAddress, unsigned int ipAddress);
+
+  bool EnableWebServer(const std::string& macAddress, unsigned int ipAddress);
+
  bool Reboot(const std::string& macAddress, unsigned int ipAddress);

+  bool GetStatus(const std::string& macAddress, unsigned int ipAddress);
+
  std::future<int>* GetFuture(const std::string& macAddress);
  void DestroyFuture(const std::string& macAddress);

-  /**
-   * Returns the state of the deploy session.
-   */
-  Status GetStatus() const;
+  std::future<DeviceStatus>* GetStatusFuture(const std::string& macAddress);
+  void DestroyStatusFuture(const std::string& macAddress);

 private:
  // Logger reference where log messages will be sent.
--- a/roborioteamnumbersetter/src/main/native/cpp/SshSession.cpp
+++ b/roborioteamnumbersetter/src/main/native/cpp/SshSession.cpp
@@ -86,13 +86,28 @@ void SshSession::Execute(std::string_view cmd) {
    ssh_channel_free(channel);
    throw SshException(ssh_get_error(m_session));
  }
-  INFO("{}", cmd);
+  INFO("{} {}", ssh_channel_get_exit_status(channel), cmd);

  // Log output.
  char buf[512];
  int read = ssh_channel_read(channel, buf, sizeof(buf), 0);
  if (read != 0) {
-    INFO("{}", cmd);
+    if (read < static_cast<int>(sizeof(buf) / sizeof(buf[0]))) {
+      buf[read] = 0;
+    } else {
+      buf[(sizeof(buf) / sizeof(buf[0])) - 1] = 0;
+    }
+    INFO("stdout: {} {}", read, buf);
+  }
+
+  read = ssh_channel_read(channel, buf, sizeof(buf), 1);
+  if (read != 0) {
+    if (read < static_cast<int>(sizeof(buf) / sizeof(buf[0]))) {
+      buf[read] = 0;
+    } else {
+      buf[(sizeof(buf) / sizeof(buf[0])) - 1] = 0;
+    }
+    INFO("stderr: {} {}", read, buf);
  }

  // Close and free channel.
@@ -100,6 +115,64 @@ void SshSession::Execute(std::string_view cmd) {
  ssh_channel_free(channel);
 }

+std::string SshSession::ExecuteResult(std::string_view cmd, int* exitStatus) {
+  // Allocate a new channel.
+  ssh_channel channel = ssh_channel_new(m_session);
+  if (!channel) {
+    throw SshException(ssh_get_error(m_session));
+  }
+
+  // Open the channel.
+  int rc = ssh_channel_open_session(channel);
+  if (rc != SSH_OK) {
+    throw SshException(ssh_get_error(m_session));
+  }
+
+  // Execute the command.
+  std::string command{cmd};
+  rc = ssh_channel_request_exec(channel, command.c_str());
+  if (rc != SSH_OK) {
+    ssh_channel_close(channel);
+    ssh_channel_free(channel);
+    throw SshException(ssh_get_error(m_session));
+  }
+  INFO("{} {}", ssh_channel_get_exit_status(channel), cmd);
+
+  std::string result;
+  if (exitStatus) {
+    *exitStatus = ssh_channel_get_exit_status(channel);
+  }
+
+  // Log output.
+  char buf[512];
+  int read = ssh_channel_read(channel, buf, sizeof(buf), 0);
+  if (read != 0) {
+    if (read < static_cast<int>(sizeof(buf) / sizeof(buf[0]))) {
+      buf[read] = 0;
+    } else {
+      buf[(sizeof(buf) / sizeof(buf[0])) - 1] = 0;
+    }
+    result = buf;
+    INFO("stdout: {} {}", read, buf);
+  }
+
+  read = ssh_channel_read(channel, buf, sizeof(buf), 1);
+  if (read != 0) {
+    if (read < static_cast<int>(sizeof(buf) / sizeof(buf[0]))) {
+      buf[read] = 0;
+    } else {
+      buf[(sizeof(buf) / sizeof(buf[0])) - 1] = 0;
+    }
+    INFO("stderr: {} {}", read, buf);
+  }
+
+  // Close and free channel.
+  ssh_channel_close(channel);
+  ssh_channel_free(channel);
+
+  return result;
+}
+
 void SshSession::Put(std::string_view path, std::string_view contents) {
  // Allocate the SFTP session.
  sftp_session sftp = sftp_new(m_session);
--- a/roborioteamnumbersetter/src/main/native/cpp/SshSession.h
+++ b/roborioteamnumbersetter/src/main/native/cpp/SshSession.h
@@ -58,6 +58,8 @@ class SshSession {
   */
  void Execute(std::string_view cmd);

+  std::string ExecuteResult(std::string_view cmd, int* exitStatus);
+
  /**
   * Puts a file on the server using SFTP.
   *