From 1bef73ed197364ac5a0f2a057fa57a2cc233ad70 Mon Sep 17 00:00:00 2001
From: Daniel Trugman <dtrugman@gmail.com>
Date: Thu, 14 Dec 2023 21:26:46 +0000
Subject: [PATCH] Add netstat sample application

---
 README.md               |  2 +
 sample/sample.cpp       | 20 ++++++---
 sample/tool.hpp         |  1 +
 sample/tool_netstat.cpp | 92 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 109 insertions(+), 6 deletions(-)
 create mode 100644 sample/tool_netstat.cpp

diff --git a/README.md b/README.md
index cea3cc5..8478182 100644
--- a/README.md
+++ b/README.md
@@ -91,6 +91,8 @@ How does that affect `pfs`?
 
 The directory `sample` contains a full blown application that calls all(!) the supported APIs and prints all the information gathered. When compiling the library, the sample applications is compiled as well.
 
+You can find a basic implementation of `netstat` (see `sample/tool_netstat.cpp`) and `lsmod` (see `sample/tool_lsmod.cpp`) that you can easily reuse in your projects.
+
 Anyway, here are some cool (and concise) examples:
 
 **Example 1:** Iterater over all the loaded unsigned or out-of-tree kernel modules
diff --git a/sample/sample.cpp b/sample/sample.cpp
index 751b200..6c35ca5 100644
--- a/sample/sample.cpp
+++ b/sample/sample.cpp
@@ -37,12 +37,20 @@ int main(int argc, char** argv)
 
         // clang-format off
         auto commands = std::vector<command>{
-            {command("system", "", "Enumerate system-wide information", enum_system)},
-            {command("net", "", "Enumerate network information", enum_net)},
-            {command("tasks", "[task-id]...", "Enumerate running tasks", enum_tasks)},
-            {command("fds", "[task-id]...", "Enumerate fds for a specific task", enum_fds)},
-            {command("lsmod", "[filter]", "Enumerate all loaded modules that match the filter", tool_lsmod)},
-            {command("blocks", "[block-name]...", "Enumerate block devices", enum_blocks)},
+            {command("system", "",
+                    "Enumerate system-wide information", enum_system)},
+            {command("net", "",
+                    "Enumerate network information", enum_net)},
+            {command("tasks", "[task-id]...",
+                    "Enumerate running tasks", enum_tasks)},
+            {command("fds", "[task-id]...",
+                    "Enumerate fds for a specific task", enum_fds)},
+            {command("blocks", "[block-name]...",
+                    "Enumerate block devices", enum_blocks)},
+            {command("lsmod", "[filter]",
+                    "Enumerate all loaded modules that match the filter", tool_lsmod)},
+            {command("netstat", "(tcp|udp) [task-id]...",
+                    "Enumerate all sockets of said type for tasks", tool_netstat)},
         };
         // clang-format on
 
diff --git a/sample/tool.hpp b/sample/tool.hpp
index 67aa16e..94b0b6c 100644
--- a/sample/tool.hpp
+++ b/sample/tool.hpp
@@ -18,5 +18,6 @@
 #define SAMPLE_TOOL_HPP
 
 int tool_lsmod(std::vector<std::string>&& args);
+int tool_netstat(std::vector<std::string>&& args);
 
 #endif // SAMPLE_TOOL_HPP
diff --git a/sample/tool_netstat.cpp b/sample/tool_netstat.cpp
new file mode 100644
index 0000000..e73411e
--- /dev/null
+++ b/sample/tool_netstat.cpp
@@ -0,0 +1,92 @@
+/*
+ *  Copyright 2020-present Daniel Trugman
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#include <regex>
+
+#include "format.hpp"
+#include "log.hpp"
+#include "tool.hpp"
+
+static const std::string TCP("tcp");
+static const std::string UDP("udp");
+
+void task_netstat(const pfs::task& task, const std::string& type)
+{
+    LOG("=========================================================");
+    LOG("Netstat for task ID[" << task.id() << "]");
+    LOG("=========================================================");
+
+    auto inodes = task.get_fds_inodes();
+    pfs::net::net_socket_filter filter = [&inodes](const pfs::net_socket& sock){
+        return inodes.find(sock.inode) == inodes.end();
+    };
+
+    auto net = task.get_net();
+    std::vector<pfs::net_socket> sockets;
+    if (type == TCP)
+    {
+        sockets = net.get_tcp(filter);
+    }
+    else if (type == UDP)
+    {
+        sockets = net.get_udp(filter);
+    }
+
+    print(sockets);
+}
+
+int tool_netstat(std::vector<std::string>&& args)
+{
+    if (args.size() < 1)
+    {
+        return -EINVAL;
+    }
+
+    std::string type(args[0]);
+    if (type != TCP && type != UDP)
+    {
+        return -EINVAL;
+    }
+
+    try
+    {
+        pfs::procfs pfs;
+
+        if (args.size() > 1)
+        {
+            for (unsigned arg = 1; arg < args.size(); ++arg)
+            {
+                auto id = std::stoi(args[arg]);
+                auto task = pfs.get_task(id);
+                task_netstat(task, type);
+            }
+        }
+        else
+        {
+            for (auto& task : pfs.get_processes())
+            {
+                task_netstat(task, type);
+            }
+        }
+    }
+    catch (const std::runtime_error& ex)
+    {
+        LOG("Error when printing netstat:");
+        LOG(TAB << ex.what());
+    }
+
+    return 0;
+}