(优化) 增加文件读写加密功能

src/model2mem.cpp

@@ -1,46 +1,56 @@
 #include <iostream>
 #include <fstream>
 #include <vector>
+#include <string>
 using namespace std;
 
-static int write_file(const std::string & fname, std::vector<char>& buf)
+// 加密密钥
-{
+string ENCRYPTION_KEY = "";
-	std::ofstream fs(fname, std::ios::binary | std::ios::out);
+
-	if(!fs.good())
+static int write_file(const std::string & fname, std::vector<char>& buf){
-	{
+    std::ofstream fs(fname, std::ios::binary | std::ios::out);
-		std::cerr<<fname<<" does not exist"<<std::endl;
+    if(!fs.good()){
-		return -1;
+        std::cerr<<fname<<" does not exist"<<std::endl;
-	}
+        return -1;
-    for(char c : buf) {
+    }
+    // 加密数据
+    size_t keyIndex = 0;
+    for(char &c : buf) {
+        c ^= ENCRYPTION_KEY[keyIndex];
+        keyIndex = (keyIndex + 1) % ENCRYPTION_KEY.length();
         fs << (int)((unsigned char)c) << ",";
     }
     fs.close();
-	return 0;
+    return 0;
 }
 
-static int load_file(const std::string & fname, std::vector<char>& buf)
+static int load_file(const std::string & fname, std::vector<char>& buf){
-{
+    std::ifstream fs(fname, std::ios::binary | std::ios::in);
-	std::ifstream fs(fname, std::ios::binary | std::ios::in);
+    if(!fs.good()){
-	if(!fs.good())
+        std::cerr<<fname<<" does not exist"<<std::endl;
-	{
+        return -1;
-		std::cerr<<fname<<" does not exist"<<std::endl;
+    }
-		return -1;
-	}
 
-	fs.seekg(0, std::ios::end);
+    fs.seekg(0, std::ios::end);
-	int fsize=fs.tellg();
+    int fsize=fs.tellg();
 
-	fs.seekg(0, std::ios::beg);
+    fs.seekg(0, std::ios::beg);
-	buf.resize(fsize);
+    buf.resize(fsize);
-	fs.read(buf.data(),fsize);
+    fs.read(buf.data(),fsize);
-	fs.close();
+    fs.close();
-	return 0;
+
+    // 解密数据
+    size_t keyIndex = 0;
+    for(char &c : buf) {
+        c ^= ENCRYPTION_KEY[keyIndex];
+        keyIndex = (keyIndex + 1) % ENCRYPTION_KEY.length();
+    }
+    return 0;
 }
 
-int main(int argc, char** argv)
+int main(int argc, char** argv){
-{
     if(argc < 3) {
-        cout << "please input the model file and mem file name" << endl;
+        cout << "model2mem [input] [output] <pwd>" << endl;
         return -1;
     }
 
@@ -50,6 +60,24 @@ int main(int argc, char** argv)
         cout << "do not read any content" << endl;
         return -1;
     }
-    write_file(argv[2], content);
+
-	return 0;
+    // 如果提供了密码，则进行加密操作
+    if(argc == 4) {
+        ENCRYPTION_KEY = argv[3];
+        write_file(argv[2], content);
+    } else {
+        // 否则，直接写入文件
+        std::ofstream fs(argv[2], std::ios::binary | std::ios::out);
+        if(!fs.good())
+        {
+            std::cerr<<argv[2]<<" does not exist"<<std::endl;
+            return -1;
+        }
+        for(char c : content) {
+            fs << (int)((unsigned char)c) << ",";
+        }
+        fs.close();
+    }
+
+    return 0;
 }