Make receiver work on sub directories

1 files changed, 40 insertions, 13 deletions

diff --git a/receiver.cpp b/receiver.cpp
index a3a0dc9..6be0f8a 100644
--- a/receiver.cpp
+++ b/receiver.cpp
@@ -7,6 +7,8 @@
 #include <fcntl.h>
 #include <unistd.h>
 #include <sys/socket.h>
+#include <sys/stat.h>
+#include <limits.h>
 #include <netinet/in.h>
 
 #include <mbedtls/pk.h>
@@ -78,15 +80,16 @@ static time_t now() {
 
 // Constants
 enum { SESSION_ID_LENGTH = 8, AES_KEY_LENGTH = 16, GCM_IV_LENGTH = 16, GCM_TAG_LENGTH = 16, MIN_PACKET_SIZE = 40 };
-enum { FILENAME_LENGTH = 73, SIDOFFS = 20, KEYOFFS = 37 };
+enum { DIRNAME_LENGTH = 10, FILENAME_LENGTH = 73, SIDOFFS = 20, KEYOFFS = 37 };
 
 class Session {
 public:
-    Session(uint64_t session_id, uint8_t key[AES_KEY_LENGTH], const std::string &filename) :
-            _session_id(session_id), _filename(filename) {
+    Session(uint64_t session_id, uint8_t key[AES_KEY_LENGTH], const std::string &dirname, const std::string &filename) :
+            _session_id(session_id), _dirname(dirname), _filename(filename) {
         memcpy(_key, key, AES_KEY_LENGTH);
         mbedtls_gcm_init(&_ctx);
         mbedtls_gcm_setkey(&_ctx, MBEDTLS_CIPHER_ID_AES, _key, 8 * AES_KEY_LENGTH);
+//        std::cerr << "Imported file " << _filename << " in dir " << _dirname << std::endl;
     }
 
     Session(uint64_t session_id, uint8_t key[AES_KEY_LENGTH]) : _session_id(session_id) {
@@ -98,7 +101,7 @@ public:
         char tprefix[32];
         time_t t = time(NULL);
         struct tm * jetzt = localtime(&t);
-        size_t nlen = strftime(tprefix, sizeof(tprefix), "%F-%H-%M-%S", jetzt);
+        strftime(tprefix, sizeof(tprefix), "%F-%H-%M-%S", jetzt);
 
         // Dump key
         char hexkey[2*AES_KEY_LENGTH + 1];
@@ -107,18 +110,24 @@ public:
 
         // Glue together serialisation
         char filename[FILENAME_LENGTH + 1];
-        snprintf(filename, sizeof(filename), "%s-%016" PRIx64 "-%s.log", tprefix, _session_id, hexkey);
+        size_t nlen = snprintf(filename, sizeof(filename), "%s-%016" PRIx64 "-%s.log", tprefix, _session_id, hexkey);
+        _filename = std::string(filename, filename + nlen);
 
-        // Touch file to save session_id and key
-        close(open(filename, O_WRONLY|O_CREAT, 0755));
+        _dirname = std::string(filename, DIRNAME_LENGTH);
 
-        _filename = std::string(filename, filename + FILENAME_LENGTH);
+        // Touch file to save session_id and key
+        close(ensure_file());
     }
 
     ~Session() {
         mbedtls_gcm_free(&_ctx);
     }
 
+    int ensure_file() {
+        mkdir(_dirname.c_str(), 0755);
+        return open((_dirname + "/" + _filename).c_str(), O_WRONLY | O_APPEND | O_CREAT, 0755);
+    }
+
     void write_log(const uint8_t *packet, size_t len) {
         // First check if the packet holds enough space for session id, iv and at least one gcm block
         if (len < GCM_IV_LENGTH + GCM_TAG_LENGTH) {
@@ -140,7 +149,7 @@ public:
 
         // Create output file if it doesn't exist
         if (_fd < 0)
-            _fd = ::open(_filename.c_str(), O_WRONLY | O_APPEND | O_CREAT, 0755);
+            _fd = ensure_file();
         if (_fd < 0) {
             std::cerr << "Error: Can't create file " << _filename << " for session " << std::hex << _session_id;
             return;
@@ -179,6 +188,7 @@ private:
     uint8_t                 _key[AES_KEY_LENGTH];
     int                     _fd = -1;
     time_t                  _last_access = 0;
+    std::string             _dirname;
     std::string             _filename;
     std::set<std::string>   _used_ivs;
     mbedtls_gcm_context     _ctx;
@@ -194,10 +204,10 @@ static uint8_t hex2nyble(char c)
         : 0;
 }
 
-static void import_sessions(const char *dirname) {
-    DIR * dirp = opendir(dirname);
+static void import_directory(const char *path, std::string dirname) {
+    DIR * dirp = opendir(path);
     if (!dirp)
-        errx(-1, "Fatal: Can't open dir %s\n", dirname);
+        errx(-1, "Fatal: Can't open dir %s\n", path);
 
     regex_t regex;
     if (regcomp(&regex, "^[[:digit:]]{4}-[[:digit:]][[:digit:]]-[[:digit:]][[:digit:]]-[[:digit:]][[:digit:]]-[[:digit:]][[:digit:]]-"
@@ -223,12 +233,29 @@ static void import_sessions(const char *dirname) {
         const char * hexkey = filename.c_str() + KEYOFFS;
         for (int i=0; i<16; ++i)
             aeskey[i] = (hex2nyble(hexkey[2*i]) << 4 ) | hex2nyble(hexkey[2*i+1]);
-        g_sessions[session_id] = std::make_unique<Session>(session_id, aeskey, filename);
+        g_sessions[session_id] = std::make_unique<Session>(session_id, aeskey, dirname, filename);
     }
     closedir(dirp);
     regfree(&regex);
 }
 
+static void import_sessions(const char *root_dir) {
+    char dirpath[PATH_MAX];
+    DIR * dirp = opendir(root_dir);
+    if (!dirp)
+        errx(-1, "Fatal: Can't open dir %s\n", root_dir);
+
+    struct dirent * entry;
+    while ((entry = readdir(dirp)) != NULL) {
+        if (entry->d_type != DT_DIR || entry->d_namlen != DIRNAME_LENGTH)
+            continue;
+        std::string dirname(entry->d_name, entry->d_name + entry->d_namlen);
+        snprintf(dirpath, PATH_MAX, "%s/%*s", root_dir, DIRNAME_LENGTH, entry->d_name);
+        import_directory(dirpath, dirname);
+    }
+    closedir(dirp);
+}
+
 int main() {
     mbedtls_ctr_drbg_context ctr_drbg;
     mbedtls_entropy_context entropy;