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

src/model2mem.cpp

@@ -1,28 +1,32 @@
 #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 = "";
+
+static int write_file(const std::string & fname, std::vector<char>& buf){
     std::ofstream fs(fname, std::ios::binary | std::ios::out);
-	if(!fs.good())
+    if(!fs.good()){
-	{
         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;
 }
 
-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);
-	if(!fs.good())
+    if(!fs.good()){
-	{
         std::cerr<<fname<<" does not exist"<<std::endl;
         return -1;
     }
@@ -34,13 +38,19 @@ static int load_file(const std::string & fname, std::vector<char>& buf)
     buf.resize(fsize);
     fs.read(buf.data(),fsize);
     fs.close();
+
+    // 解密数据
+    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;
     }
+
+    // 如果提供了密码，则进行加密操作
+    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;
 }