Feature: Implement user management feature with role support

This commit is contained in:
yuanyuanxiang
2026-04-24 12:19:15 +02:00
parent ac14073921
commit 655b1934a4
3 changed files with 617 additions and 3 deletions

View File

@@ -9,6 +9,8 @@
#include "SimpleWebSocket.h"
#include "common/commands.h"
#include <filesystem>
#include <fstream>
#include <shlobj.h>
// Algorithm constants (same as ScreenSpyDlg.cpp)
#define ALGORITHM_H264 2
@@ -20,6 +22,16 @@ static std::map<void*, std::string> s_ClientNonces;
static std::mutex s_NonceMutex;
static std::atomic<bool> s_bShuttingDown{false}; // Prevents access during static destruction
// Generate random salt (16 hex chars) - thread-safe
static std::string GenerateSalt() {
static std::random_device rd;
static std::mt19937_64 gen(rd());
std::uniform_int_distribution<uint64_t> dis;
char buf[17];
snprintf(buf, sizeof(buf), "%016llX", dis(gen));
return std::string(buf);
}
// Generate random nonce (32 hex chars) - thread-safe
static std::string GenerateNonce() {
if (s_bShuttingDown) return "";
@@ -112,11 +124,22 @@ CWebService::CWebService()
m_PayloadsDir = (exeDir / "Payloads").string();
std::error_code ec;
std::filesystem::create_directories(m_PayloadsDir, ec);
// Initialize config directory (same as YAMA.db location)
char appdata_path[MAX_PATH];
if (SUCCEEDED(SHGetFolderPathA(NULL, CSIDL_APPDATA, NULL, 0, appdata_path))) {
m_ConfigDir = std::string(appdata_path) + "\\" BRAND_DATA_FOLDER "\\";
} else {
m_ConfigDir = ".\\";
}
std::filesystem::create_directories(m_ConfigDir, ec);
}
void CWebService::SetAdminPassword(const std::string& password) {
std::lock_guard<std::mutex> lock(m_UsersMutex);
m_Users.clear();
// Admin user is built-in, always first
WebUser admin;
admin.username = "admin";
admin.salt = ""; // Not used with challenge-response auth
@@ -125,6 +148,9 @@ void CWebService::SetAdminPassword(const std::string& password) {
m_Users.push_back(admin);
Mprintf("[WebService] Admin password configured\n");
// Load additional users from file (non-admin users)
LoadUsers();
}
CWebService::~CWebService() {
@@ -329,6 +355,14 @@ void CWebService::ServerThread(int port) {
HandleKey(ws_ptr, msg);
} else if (cmd == "rdp_reset") {
HandleRdpReset(ws_ptr, token);
} else if (cmd == "get_salt") {
HandleGetSalt(ws_ptr, msg);
} else if (cmd == "create_user") {
HandleCreateUser(ws_ptr, msg);
} else if (cmd == "delete_user") {
HandleDeleteUser(ws_ptr, msg);
} else if (cmd == "list_users") {
HandleListUsers(ws_ptr, token);
}
}
});
@@ -480,6 +514,51 @@ void CWebService::HandleLogin(void* ws_ptr, const std::string& msg, const std::s
SendText(ws_ptr, Json::writeString(builder, res));
}
void CWebService::HandleGetSalt(void* ws_ptr, const std::string& msg) {
Json::Value root;
Json::Reader reader;
if (!reader.parse(msg, root)) {
SendText(ws_ptr, BuildJsonResponse("salt", false, "Invalid JSON"));
return;
}
std::string username = root.get("username", "").asString();
if (username.empty()) {
SendText(ws_ptr, BuildJsonResponse("salt", false, "Username required"));
return;
}
// Find user and get salt
std::string salt = "";
bool userFound = false;
{
std::lock_guard<std::mutex> lock(m_UsersMutex);
for (const auto& u : m_Users) {
if (u.username == username) {
salt = u.salt;
userFound = true;
break;
}
}
}
// For security: if user doesn't exist, generate a fake deterministic salt
// This prevents username enumeration attacks
// Note: Admin has empty salt, so we must check userFound, not salt.empty()
if (!userFound) {
// Generate deterministic fake salt from username (won't match any real password)
salt = WSAuth::ComputeSHA256("fake_salt_prefix_" + username).substr(0, 16);
}
Json::Value res;
res["cmd"] = "salt";
res["ok"] = true;
res["salt"] = salt;
Json::StreamWriterBuilder builder;
builder["indentation"] = "";
SendText(ws_ptr, Json::writeString(builder, res));
}
void CWebService::HandleGetDevices(void* ws_ptr, const std::string& token) {
std::string username, role;
if (!ValidateToken(token, username, role)) {
@@ -837,6 +916,111 @@ void CWebService::HandleRdpReset(void* ws_ptr, const std::string& token) {
}
}
//////////////////////////////////////////////////////////////////////////
// User Management Handlers
//////////////////////////////////////////////////////////////////////////
void CWebService::HandleCreateUser(void* ws_ptr, const std::string& msg) {
Json::Value root;
Json::Reader reader;
if (!reader.parse(msg, root)) {
SendText(ws_ptr, BuildJsonResponse("create_user_result", false, "Invalid JSON"));
return;
}
std::string token = root.get("token", "").asString();
std::string username, role;
if (!ValidateToken(token, username, role)) {
SendText(ws_ptr, BuildJsonResponse("create_user_result", false, "Invalid token"));
return;
}
// Only admin can create users
if (role != "admin") {
SendText(ws_ptr, BuildJsonResponse("create_user_result", false, "Permission denied"));
return;
}
std::string newUsername = root.get("username", "").asString();
std::string newPassword = root.get("password", "").asString();
std::string newRole = root.get("role", "viewer").asString();
if (newUsername.empty() || newPassword.empty()) {
SendText(ws_ptr, BuildJsonResponse("create_user_result", false, "Username and password required"));
return;
}
if (CreateUser(newUsername, newPassword, newRole)) {
SendText(ws_ptr, BuildJsonResponse("create_user_result", true));
} else {
SendText(ws_ptr, BuildJsonResponse("create_user_result", false, "Failed to create user (may already exist)"));
}
}
void CWebService::HandleDeleteUser(void* ws_ptr, const std::string& msg) {
Json::Value root;
Json::Reader reader;
if (!reader.parse(msg, root)) {
SendText(ws_ptr, BuildJsonResponse("delete_user_result", false, "Invalid JSON"));
return;
}
std::string token = root.get("token", "").asString();
std::string username, role;
if (!ValidateToken(token, username, role)) {
SendText(ws_ptr, BuildJsonResponse("delete_user_result", false, "Invalid token"));
return;
}
// Only admin can delete users
if (role != "admin") {
SendText(ws_ptr, BuildJsonResponse("delete_user_result", false, "Permission denied"));
return;
}
std::string targetUsername = root.get("username", "").asString();
if (DeleteUser(targetUsername)) {
SendText(ws_ptr, BuildJsonResponse("delete_user_result", true));
} else {
SendText(ws_ptr, BuildJsonResponse("delete_user_result", false, "Failed to delete user"));
}
}
void CWebService::HandleListUsers(void* ws_ptr, const std::string& token) {
std::string username, role;
if (!ValidateToken(token, username, role)) {
SendText(ws_ptr, BuildJsonResponse("list_users_result", false, "Invalid token"));
return;
}
// Only admin can list users
if (role != "admin") {
SendText(ws_ptr, BuildJsonResponse("list_users_result", false, "Permission denied"));
return;
}
auto users = ListUsers();
Json::Value res;
res["cmd"] = "list_users_result";
res["ok"] = true;
Json::Value usersArray(Json::arrayValue);
for (const auto& u : users) {
Json::Value user;
user["username"] = u.first;
user["role"] = u.second;
usersArray.append(user);
}
res["users"] = usersArray;
Json::StreamWriterBuilder builder;
builder["indentation"] = "";
std::string json = Json::writeString(builder, res);
SendText(ws_ptr, json);
}
//////////////////////////////////////////////////////////////////////////
// Token Management (delegated to WebServiceAuth module)
//////////////////////////////////////////////////////////////////////////
@@ -936,6 +1120,155 @@ std::string CWebService::ComputeHash(const std::string& input) {
return WSAuth::ComputeSHA256(input);
}
//////////////////////////////////////////////////////////////////////////
// User Management
//////////////////////////////////////////////////////////////////////////
std::string CWebService::GetUsersFilePath() {
return m_ConfigDir + "users.json";
}
void CWebService::LoadUsers() {
// Note: m_UsersMutex should already be held by caller (SetAdminPassword)
// Load additional users from users.json (admin user is already in m_Users)
std::string path = GetUsersFilePath();
std::ifstream file(path);
if (!file.is_open()) {
Mprintf("[WebService] No users.json found, using admin only\n");
return;
}
try {
Json::Value root;
Json::CharReaderBuilder builder;
std::string errors;
if (!Json::parseFromStream(builder, file, &root, &errors)) {
Mprintf("[WebService] Failed to parse users.json: %s\n", errors.c_str());
return;
}
const Json::Value& users = root["users"];
int loaded = 0;
for (const auto& u : users) {
std::string username = u.get("username", "").asString();
// Skip admin user (it's built-in with master password)
if (username.empty() || username == "admin") continue;
WebUser user;
user.username = username;
user.password_hash = u.get("password_hash", "").asString();
user.salt = u.get("salt", "").asString();
user.role = u.get("role", "viewer").asString();
if (!user.password_hash.empty()) {
m_Users.push_back(user);
loaded++;
}
}
Mprintf("[WebService] Loaded %d additional users from users.json\n", loaded);
} catch (const std::exception& e) {
Mprintf("[WebService] Error loading users.json: %s\n", e.what());
}
}
void CWebService::SaveUsers() {
// Save non-admin users to users.json
std::lock_guard<std::mutex> lock(m_UsersMutex);
Json::Value root;
Json::Value users(Json::arrayValue);
for (const auto& u : m_Users) {
// Skip admin user (it uses master password, not stored in file)
if (u.username == "admin") continue;
Json::Value user;
user["username"] = u.username;
user["password_hash"] = u.password_hash;
user["salt"] = u.salt;
user["role"] = u.role;
users.append(user);
}
root["users"] = users;
std::string path = GetUsersFilePath();
std::ofstream file(path);
if (!file.is_open()) {
Mprintf("[WebService] Failed to open users.json for writing\n");
return;
}
Json::StreamWriterBuilder builder;
builder["indentation"] = " ";
std::unique_ptr<Json::StreamWriter> writer(builder.newStreamWriter());
writer->write(root, &file);
Mprintf("[WebService] Saved %d users to users.json\n", (int)users.size());
}
bool CWebService::CreateUser(const std::string& username, const std::string& password, const std::string& role) {
if (username.empty() || password.empty()) return false;
if (username == "admin") return false; // Cannot create user named "admin"
if (role != "admin" && role != "viewer") return false;
{
std::lock_guard<std::mutex> lock(m_UsersMutex);
// Check if user already exists
for (const auto& u : m_Users) {
if (u.username == username) return false;
}
// Generate salt and hash password with salt
WebUser user;
user.username = username;
user.salt = GenerateSalt();
user.password_hash = WSAuth::ComputeSHA256(password + user.salt);
user.role = role;
m_Users.push_back(user);
Mprintf("[WebService] Created user: %s (role: %s)\n", username.c_str(), role.c_str());
}
// Save to file (outside lock scope since SaveUsers acquires its own lock)
SaveUsers();
return true;
}
bool CWebService::DeleteUser(const std::string& username) {
if (username.empty() || username == "admin") return false;
bool deleted = false;
{
std::lock_guard<std::mutex> lock(m_UsersMutex);
for (auto it = m_Users.begin(); it != m_Users.end(); ++it) {
if (it->username == username) {
m_Users.erase(it);
Mprintf("[WebService] Deleted user: %s\n", username.c_str());
deleted = true;
break;
}
}
}
if (deleted) {
SaveUsers();
}
return deleted;
}
std::vector<std::pair<std::string, std::string>> CWebService::ListUsers() {
std::lock_guard<std::mutex> lock(m_UsersMutex);
std::vector<std::pair<std::string, std::string>> result;
for (const auto& u : m_Users) {
result.push_back({u.username, u.role});
}
return result;
}
//////////////////////////////////////////////////////////////////////////
// JSON Helpers
//////////////////////////////////////////////////////////////////////////