Skip to content

Commit df891eb

Browse files
Update cpp
1 parent 7dfba9f commit df891eb

26 files changed

Lines changed: 1184 additions & 123 deletions

‎README.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@
6060

6161
## Getting Started
6262

63-
#### Run **heiDGAF** using Docker Compose:
63+
#### Run **HAMSTRING** using Docker Compose:
6464
```sh
6565
HOST_IP=127.0.0.1 docker compose -f docker/docker-compose.yml --profile prod up
6666
```
@@ -98,7 +98,7 @@ python src/inspector/inspector.py
9898

9999
### Configuration
100100

101-
To configure **heiDGAF** according to your needs, use the provided `config.yaml`.
101+
To configure **HAMSTRING** according to your needs, use the provided `config.yaml`.
102102

103103
The most relevant settings are related to your specific log line format, the model you want to use, and
104104
possibly infrastructure.
@@ -289,7 +289,7 @@ This will create a `rules.txt` file containing the innards of the model, explain
289289
> We support custom schemes.
290290

291291
Depending on your data and usecase, you can customize the data scheme to fit your needs.
292-
The below configuration is part of the [main configuration file](./config.yaml) which is detailed in our [documentation](https://heidgaf.readthedocs.io/en/latest/usage.html#id2)
292+
The below configuration is part of the [main configuration file](./config.yaml) which is detailed in our [documentation](https://HAMSTRING.readthedocs.io/en/latest/usage.html#id2)
293293

294294
```yml
295295
loglines:

‎cpp/CMakeLists.txt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,11 +31,11 @@ set(CMAKE_CXX_FLAGS_DEBUG "-g -O0")
3131
# Find packages
3232
find_package(yaml-cpp CONFIG REQUIRED)
3333
find_package(RdKafka CONFIG REQUIRED)
34-
# find_package(clickhouse-cpp CONFIG REQUIRED) # TODO: Fix clickhouse-cpp
34+
# find_package(clickhouse-cpp CONFIG REQUIRED) # No CMake config provided by vcpkg
3535
find_package(Boost REQUIRED COMPONENTS system thread)
3636
find_package(spdlog CONFIG REQUIRED)
3737
find_package(fmt CONFIG REQUIRED)
38-
# find_package(onnxruntime CONFIG REQUIRED) # TODO: Add ONNX support later
38+
# find_package(onnxruntime CONFIG REQUIRED) # Optional for Detector module
3939
find_package(nlohmann_json CONFIG REQUIRED)
4040
find_package(OpenSSL REQUIRED)
4141

‎cpp/examples/CMakeLists.txt‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
# Example executable
2+
add_executable(demo
3+
demo.cpp
4+
)
5+
6+
target_link_libraries(demo
7+
PRIVATE
8+
hamstring_base
9+
hamstring_detector
10+
)

‎cpp/examples/demo.cpp‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
#include "hamstring/base/logger.hpp"
2+
#include "hamstring/base/utils.hpp"
3+
#include "hamstring/config/config.hpp"
4+
#include "hamstring/detector/feature_extractor.hpp"
5+
#include <iostream>
6+
7+
using namespace hamstring;
8+
9+
int main(int argc, char **argv) {
10+
// Initialize logger
11+
base::Logger::initialize(true); // debug mode
12+
auto logger = base::Logger::get_logger("example");
13+
14+
logger->info("HAMSTRING C++ Example");
15+
logger->info("=====================");
16+
17+
// Load configuration
18+
std::string config_path = (argc > 1) ? argv[1] : "../../config.yaml";
19+
logger->info("Loading configuration from: {}", config_path);
20+
21+
try {
22+
auto config = config::Config::load_from_file(config_path);
23+
24+
logger->info("Configuration loaded successfully");
25+
logger->info("Number of collectors: {}",
26+
config->pipeline.collectors.size());
27+
logger->info("Number of detectors: {}", config->pipeline.detectors.size());
28+
logger->info("Kafka brokers: {}", config->environment.kafka_brokers.size());
29+
30+
// Show Kafka bootstrap servers
31+
std::string bootstrap = config->environment.get_kafka_bootstrap_servers();
32+
logger->info("Kafka bootstrap servers: {}", bootstrap);
33+
34+
} catch (const std::exception &e) {
35+
logger->error("Failed to load configuration: {}", e.what());
36+
logger->warn("Continuing with feature extraction demo...");
37+
}
38+
39+
// Demonstrate feature extraction
40+
logger->info("");
41+
logger->info("Feature Extraction Demo");
42+
logger->info("=======================");
43+
44+
detector::FeatureExtractor extractor;
45+
46+
// Test domains
47+
std::vector<std::string> test_domains = {"google.com", "www.example.com",
48+
"xjk3n2m9pq.com", // DGA-like
49+
"mail.google.com"};
50+
51+
for (const auto &domain : test_domains) {
52+
logger->info("");
53+
logger->info("Domain: {}", domain);
54+
55+
auto features = extractor.extract(domain);
56+
57+
logger->info(" Label length: {}", features.label_length);
58+
logger->info(" Label max: {}", features.label_max);
59+
logger->info(" FQDN entropy: {:.4f}", features.fqdn_entropy);
60+
logger->info(" SLD entropy: {:.4f}", features.secondleveldomain_entropy);
61+
logger->info(" Alpha ratio: {:.4f}", features.fqdn_alpha_count);
62+
logger->info(" Numeric ratio: {:.4f}", features.fqdn_numeric_count);
63+
64+
auto vec = features.to_vector();
65+
logger->info(" Feature vector size: {}", vec.size());
66+
67+
// Check if DGA-like (high entropy)
68+
if (features.fqdn_entropy > 3.0) {
69+
logger->warn(" ⚠ High entropy - possible DGA domain!");
70+
}
71+
}
72+
73+
// Demonstrate utilities
74+
logger->info("");
75+
logger->info("Utilities Demo");
76+
logger->info("==============");
77+
78+
std::string uuid = base::utils::generate_uuid();
79+
logger->info("Generated UUID: {}", uuid);
80+
81+
std::string ip = "192.168.1.100";
82+
std::string subnet = base::utils::get_subnet_id(ip, 24);
83+
logger->info("IP: {} -> Subnet: {}", ip, subnet);
84+
85+
std::string domain = "www.example.com";
86+
std::string sld = base::utils::extract_second_level_domain(domain);
87+
std::string tld_part = base::utils::extract_third_level_domain(domain);
88+
logger->info("Domain: {}", domain);
89+
logger->info(" Second level: {}", sld);
90+
logger->info(" Third level: {}", tld_part);
91+
92+
logger->info("");
93+
logger->info("Example completed successfully!");
94+
95+
return 0;
96+
}

‎cpp/include/hamstring/base/clickhouse_sender.hpp‎

Lines changed: 33 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,17 @@
11
#pragma once
22

3-
#include <chrono>
4-
// #include <clickhouse/client.h> // TODO: Enable when clickhouse-cpp is
5-
// available
6-
#include <memory>
73
#include <string>
8-
#include <vector>
94

105
namespace hamstring {
116
namespace base {
127

13-
// ClickHouse batch sender for monitoring data
8+
/**
9+
* @brief ClickHouse stub for monitoring (logs to spdlog instead of database)
10+
*
11+
* This is a lightweight stub that provides the ClickHouse interface but
12+
* logs metrics instead of writing to a database. Useful for testing and
13+
* when ClickHouse is not available.
14+
*/
1415
class ClickHouseSender {
1516
public:
1617
ClickHouseSender(const std::string &hostname, int port = 9000,
@@ -19,56 +20,39 @@ class ClickHouseSender {
1920
const std::string &password = "");
2021
~ClickHouseSender();
2122

22-
// Execute a query
23+
// Batch tracking (logged as structured messages)
24+
void insert_batch_timestamp(const std::string &batch_id,
25+
const std::string &stage,
26+
const std::string &instance_name,
27+
const std::string &status, size_t message_count,
28+
bool is_active = true);
29+
30+
// Logline tracking
31+
void insert_logline_timestamp(const std::string &logline_id,
32+
const std::string &stage,
33+
const std::string &status,
34+
bool is_active = true);
35+
36+
// Metrics/fill levels
37+
void insert_fill_level(const std::string &stage,
38+
const std::string &entry_type, size_t entry_count);
39+
40+
// DGA detections
41+
void insert_dga_detection(const std::string &domain, double score,
42+
const std::string &batch_id,
43+
const std::string &src_ip);
44+
45+
// Generic methods
2346
void execute(const std::string &query);
24-
25-
// Insert data into a table
26-
void insert(const std::string &table, const std::vector<std::string> &columns,
27-
const std::vector<std::vector<std::string>> &rows);
28-
29-
// Check if connection is alive
3047
bool ping();
3148

3249
private:
3350
std::string hostname_;
3451
int port_;
3552
std::string database_;
36-
// std::unique_ptr<clickhouse::Client> client_; // TODO: Enable when
37-
// clickhouse-cpp is available
38-
};
39-
40-
// Batched ClickHouse sender with timeout
41-
class ClickHouseBatchSender {
42-
public:
43-
ClickHouseBatchSender(
44-
std::shared_ptr<ClickHouseSender> sender, const std::string &table,
45-
const std::vector<std::string> &columns, size_t batch_size = 50,
46-
std::chrono::seconds batch_timeout = std::chrono::seconds(2));
47-
~ClickHouseBatchSender();
48-
49-
// Add a row to the batch
50-
void add_row(const std::vector<std::string> &row);
51-
52-
// Flush the batch
53-
void flush();
54-
55-
// Start auto-flush timer
56-
void start_auto_flush();
57-
58-
// Stop auto-flush timer
59-
void stop_auto_flush();
60-
61-
private:
62-
void check_and_flush();
63-
64-
std::shared_ptr<ClickHouseSender> sender_;
65-
std::string table_;
66-
std::vector<std::string> columns_;
67-
size_t batch_size_;
68-
std::chrono::seconds batch_timeout_;
69-
std::vector<std::vector<std::string>> batch_;
70-
std::chrono::system_clock::time_point last_flush_;
71-
bool auto_flush_running_ = false;
53+
std::string user_;
54+
std::string password_;
55+
bool connected_;
7256
};
7357

7458
} // namespace base
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
#pragma once
2+
3+
#include "hamstring/detector/feature_extractor.hpp"
4+
#include <memory>
5+
#include <string>
6+
#include <vector>
7+
8+
// Forward declaration for ONNX Runtime classes to avoid exposing them in the
9+
// header
10+
namespace onnxruntime {
11+
class InferenceSession;
12+
class Env;
13+
class SessionOptions;
14+
class RunOptions;
15+
} // namespace onnxruntime
16+
17+
namespace hamstring {
18+
namespace detector {
19+
20+
class Detector {
21+
public:
22+
Detector();
23+
~Detector();
24+
25+
// Load ONNX model from file
26+
void load_model(const std::string &model_path);
27+
28+
// Predict probability of domain being DGA (0.0 - 1.0)
29+
float predict(const std::string &domain);
30+
31+
private:
32+
FeatureExtractor feature_extractor_;
33+
34+
// Pimpl idiom for ONNX Runtime objects
35+
struct Impl;
36+
std::unique_ptr<Impl> impl_;
37+
};
38+
39+
} // namespace detector
40+
} // namespace hamstring
Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
```cpp
2+
#pragma once
3+
4+
#include "hamstring/base/clickhouse_sender.hpp"
5+
#include "hamstring/base/data_classes.hpp"
6+
#include "hamstring/base/kafka_handler.hpp"
7+
#include "hamstring/base/logger.hpp"
8+
#include "hamstring/config/config.hpp"
9+
#include "hamstring/detector/detector.hpp"
10+
#include <atomic>
11+
#include <memory>
12+
#include <spdlog/spdlog.h>
13+
#include <string>
14+
#include <thread>
15+
#include <vector>
16+
17+
namespace hamstring {
18+
namespace detector {
19+
20+
class DetectorService {
21+
public:
22+
DetectorService(const std::string &name, const std::string &consume_topic,
23+
const std::string &model_path, double threshold,
24+
std::shared_ptr<config::Config> config,
25+
const std::string &bootstrap_servers,
26+
const std::string &group_id);
27+
~DetectorService();
28+
29+
void start();
30+
void stop();
31+
bool is_running() const { return running_; }
32+
33+
struct Stats {
34+
uint64_t batches_consumed = 0;
35+
uint64_t domains_scanned = 0;
36+
uint64_t domains_detected = 0;
37+
};
38+
Stats get_stats() const;
39+
40+
private:
41+
void consume_loop();
42+
void process_batch(const base::Batch &batch);
43+
44+
std::string name_;
45+
std::string consume_topic_;
46+
std::string model_path_;
47+
double threshold_;
48+
std::shared_ptr<config::Config> config_;
49+
50+
std::shared_ptr<spdlog::logger> logger_;
51+
std::unique_ptr<base::KafkaConsumeHandler> consumer_;
52+
std::shared_ptr<base::ClickHouseSender> clickhouse_;
53+
54+
// The core detector logic
55+
Detector detector_;
56+
57+
std::atomic<bool> running_{false};
58+
std::thread worker_thread_;
59+
60+
// Stats
61+
std::atomic<uint64_t> batches_consumed_{0};
62+
std::atomic<uint64_t> domains_scanned_{0};
63+
std::atomic<uint64_t> domains_detected_{0};
64+
};
65+
66+
// Factory function
67+
std::vector<std::shared_ptr<DetectorService>>
68+
create_detector_services(std::shared_ptr<config::Config> config);
69+
70+
} // namespace detector
71+
} // namespace hamstring

0 commit comments

Comments
 (0)