├── CMakeLists.txt
├── README.md
├── caffe.pb.cc
├── caffe.pb.h
├── caffe_encrypt.cpp
├── caffe_encrypt.h
└── encrypt_tool.cpp


/CMakeLists.txt:
--------------------------------------------------------------------------------
 1 | cmake_minimum_required(VERSION 3.5)
 2 | project(uniubi_license)
 3 | set(CMAKE_BUILD_TYPE Release)
 4 | set(CMAKE_CXX_STANDARD 11)
 5 | if(CMAKE_BUILD_TYPE MATCHES "None")
 6 |     set(CMAKE_BUILD_TYPE Release)
 7 | endif()
 8 | 
 9 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11")
10 | 
11 | find_package(Protobuf REQUIRED)
12 | 
13 | add_library(caffe_encrypt SHARED caffe_encrypt.cpp caffe.pb.cc)
14 | target_link_libraries(caffe_encrypt ${PROTOBUF_LIBRARY} glog)
15 | add_executable(encrypt_tool encrypt_tool.cpp)
16 | target_link_libraries(encrypt_tool caffe_encrypt)
17 | 
18 | #add_executable(test_feature test.cpp)
19 | #target_link_libraries(test_feature blobd)


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # a simple tool for encrypt and decrypt caffe model
 2 | **prequirement**:  
 3 | protobuf  
 4 | glog  
 5 | ### how to use:   
 6 | 1.build
 7 | ```
 8 | mkdir build
 9 | cd build
10 | cmake ..
11 | make
12 | ```  
13 | 2.encrypt a model  
14 | encrypt_tool has 3 args, the first one is the model prefix with its path,for example if you have a model named mnist which has mnist.prototxt and mnist.caffemodel,then the prefix will be $PATH/mnist,the second one is the password for encrypt,the last one is the output file name, for example:
15 | ```
16 | ./encrypt_tool mode_prefix passwd output_name
17 | ```  
18 | 3.decrypt a model for use  
19 | let's say we are going to init a caffe Net with encrypt file: mnist.prototxt and mnist.caffemodel, the code will be like:
20 | ```
21 | shared_ptr<Net> net_ptr;
22 | std::string pwd="my_encrypt_path";
23 | 
24 | NetParameter proto;
25 | caffe_encrypt::decrypt_proto(protopath,pwd,proto);//decrypt proto
26 | net_ptr.reset(new Net{proto,caffe::TEST});
27 | 
28 | NetParameter modelparams;
29 | caffe_encrypt::decrypt_binary(modelpath,pwd,modelparams);//decrypt params
30 | 
31 | net_ptr->CopyTrainedLayersFrom(modelparams);
32 | 
33 | then you can use the model as before
34 | ```
35 | **tips**:
36 | I am using nvcaffe, maybe have some different as the original caffe
37 | 


--------------------------------------------------------------------------------
/caffe_encrypt.cpp:
--------------------------------------------------------------------------------
  1 | //
  2 | // Created by liuhao on 9/6/18.
  3 | //
  4 | 
  5 | #include "caffe_encrypt.h"
  6 | #include <fstream>
  7 | #include <google/protobuf/text_format.h>
  8 | #include <google/protobuf/io/coded_stream.h>
  9 | #include <google/protobuf/io/zero_copy_stream_impl.h>
 10 | #include <glog/logging.h>
 11 | #include <fcntl.h>
 12 | #include <sstream>
 13 | #include "caffe.pb.h"
 14 | 
 15 | using google::protobuf::io::FileInputStream;
 16 | using google::protobuf::io::ZeroCopyInputStream;
 17 | using google::protobuf::io::CodedInputStream;
 18 | using google::protobuf::Message;
 19 | 
 20 | void caffe_encrypt::encrypt_binary(const std::string &file_name, const std::string &passwd,
 21 |                                    const std::string &output_file){
 22 |     if(passwd.empty())
 23 |         return;
 24 | 
 25 |     std::ifstream in_file(file_name,std::ios::binary);
 26 | 
 27 |     if(!in_file)
 28 |         return;
 29 | 
 30 |     std::streampos begin, end;
 31 |     begin = in_file.tellg();
 32 |     in_file.seekg(0, std::ios::end);
 33 |     end = in_file.tellg();
 34 |     std::size_t size = end - begin;
 35 |     fprintf(stdout, "engine file size: %d bytes\n", size);
 36 |     in_file.seekg(0, std::ios::beg);
 37 | 
 38 |     std::vector<unsigned char> engine_data(size);
 39 |     in_file.read((char*)engine_data.data(), size);
 40 |     in_file.close();
 41 | 
 42 |     std::ofstream out_proto(output_file, std::ios::out | std::ios::binary);
 43 |     if (!out_proto.is_open()) {
 44 |         fprintf(stderr, "fail to open file to write: %s\n", output_file.c_str());
 45 |         return;
 46 |     }
 47 | 
 48 |     //encode
 49 |     if(passwd.size()>0)
 50 |     {
 51 |         int password_len=passwd.size();
 52 |         for(int i=0;i<engine_data.size();++i)
 53 |         {
 54 |             unsigned char before_encode=engine_data[i];
 55 |             unsigned char encode_word=(before_encode)^(unsigned char)passwd[i%password_len];
 56 |             out_proto.put(encode_word);
 57 |             //outfile.write((char*)&encode_word, sizeof(char));
 58 |         }
 59 |     }
 60 | 
 61 |     out_proto.close();
 62 | }
 63 | 
 64 | void caffe_encrypt::encrypt_proto(const std::string &file_name, const std::string &passwd,
 65 |                                    const std::string &output_file){
 66 |     if(passwd.empty())
 67 |         return;
 68 | 
 69 |     std::ifstream in_file(file_name);
 70 | 
 71 |     if(!in_file)
 72 |         return;
 73 | 
 74 |     std::streampos begin, end;
 75 |     begin = in_file.tellg();
 76 |     in_file.seekg(0, std::ios::end);
 77 |     end = in_file.tellg();
 78 |     std::size_t size = end - begin;
 79 |     fprintf(stdout, "engine file size: %d bytes\n", size);
 80 |     in_file.seekg(0, std::ios::beg);
 81 | 
 82 |     std::vector<unsigned char> engine_data(size);
 83 |     in_file.read((char*)engine_data.data(), size);
 84 |     in_file.close();
 85 | 
 86 |     std::ofstream out_file(output_file,std::ios::binary);
 87 | 
 88 |     //encode
 89 |     if(passwd.size()>0)
 90 |     {
 91 |         int password_len=passwd.size();
 92 |         for(int i=0;i<engine_data.size();++i)
 93 |         {
 94 |             unsigned char before_encode=engine_data[i];
 95 |             unsigned char encode_word=(before_encode)^(unsigned char)passwd[i%password_len];
 96 |             out_file.put(encode_word);
 97 |             //outfile.write((char*)&encode_word, sizeof(char));
 98 |         }
 99 |     }
100 | 
101 |     out_file.close();
102 | }
103 | 
104 | void caffe_encrypt::encrypt(const std::string &prefix, const std::string &passwd,
105 |                             const std::string &output_name)
106 | {
107 |     if(passwd.empty())
108 |         return;
109 | 
110 |     //encode prototxt
111 |     encrypt_proto(prefix+".prototxt",passwd,output_name+".prototxt");
112 | 
113 |     //encode model
114 |     encrypt_binary(prefix+".caffemodel",passwd,output_name+".caffemodel");
115 | }
116 | 
117 | void caffe_encrypt::decrypt_binary(const std::string& file_name,const std::string& passwd,caffe::NetParameter& params)
118 | {
119 |     std::ifstream in_file(file_name,std::ios::binary);
120 | 
121 |     if(!in_file)
122 |         return;
123 | 
124 |     std::streampos begin, end;
125 |     begin = in_file.tellg();
126 |     in_file.seekg(0, std::ios::end);
127 |     end = in_file.tellg();
128 |     std::size_t size = end - begin;
129 |     fprintf(stdout, "engine file size: %d bytes\n", size);
130 |     in_file.seekg(0, std::ios::beg);
131 | 
132 |     std::string engine_data(size,'c');
133 |     in_file.read((char*)engine_data.data(), size);
134 |     in_file.close();
135 | 
136 |     int pwd_len=passwd.size();
137 | 
138 |     if(pwd_len>0)
139 |     {
140 |         //decrypt
141 |         for (size_t i = 0; i < engine_data.size(); ++i)
142 |         {
143 |             engine_data[i]=(((unsigned char)engine_data[i])^((unsigned char)passwd[i%pwd_len]));
144 |         }
145 |     }
146 | 
147 |     std::istringstream net_model(engine_data);
148 |     google::protobuf::io::IstreamInputStream* model_input=new google::protobuf::io::IstreamInputStream((std::istream*)(&net_model));
149 |     google::protobuf::io::CodedInputStream* coded_input_model=new google::protobuf::io::CodedInputStream(model_input);
150 |     coded_input_model->SetTotalBytesLimit(INT_MAX,536870912);
151 |     params.ParseFromCodedStream(coded_input_model);
152 | 
153 |     delete model_input;
154 |     delete coded_input_model;
155 | 
156 |     return;
157 | }
158 | 
159 | void caffe_encrypt::decrypt_proto(const std::string &file_name, const std::string &passwd, caffe::NetParameter& params){
160 |     std::ifstream in_file(file_name,std::ios::binary);
161 |     if(!in_file)
162 |         return;
163 | 
164 |     std::streampos begin, end;
165 |     begin = in_file.tellg();
166 |     in_file.seekg(0, std::ios::end);
167 |     end = in_file.tellg();
168 |     std::size_t size = end - begin;
169 |     fprintf(stdout, "engine file size: %d bytes\n", size);
170 |     in_file.seekg(0, std::ios::beg);
171 | 
172 |     std::string engine_data(size,'c');
173 |     in_file.read((char*)engine_data.data(), size);
174 |     in_file.close();
175 | 
176 |     int pwd_len=passwd.size();
177 | 
178 |     if(pwd_len>0)
179 |     {
180 |         //decrypt
181 |         for (size_t i = 0; i < engine_data.size(); ++i)
182 |         {
183 |             engine_data[i]=(((unsigned char)engine_data[i])^((unsigned char)passwd[i%pwd_len]));
184 |         }
185 |     }
186 | 
187 |     std::istringstream net_proto(engine_data);
188 | 
189 |     google::protobuf::io::IstreamInputStream* proto_input=new google::protobuf::io::IstreamInputStream((std::istream*)(&net_proto));
190 |     google::protobuf::io::ZeroCopyInputStream* coded_input_proto=proto_input;
191 | 
192 |     google::protobuf::TextFormat::Parse(coded_input_proto,&params);
193 | 
194 |     delete proto_input;
195 | 
196 |     return;
197 | }


--------------------------------------------------------------------------------
/caffe_encrypt.h:
--------------------------------------------------------------------------------
 1 | //
 2 | // Created by liuhao on 9/6/18.
 3 | //
 4 | 
 5 | #ifndef GRPC_DATAFORMAT_CAFFE_ENCRYPT_H
 6 | #define GRPC_DATAFORMAT_CAFFE_ENCRYPT_H
 7 | 
 8 | #include <string>
 9 | #include <google/protobuf/text_format.h>
10 | namespace caffe
11 | {
12 |     class NetParameter;
13 | }
14 | class caffe_encrypt
15 | {
16 | public:
17 |     static void encrypt(const std::string& prefix, const std::string& passwd,const std::string& output_name);
18 |     static void decrypt_proto(const std::string& file_name,const std::string& passwd,caffe::NetParameter& params);
19 |     static void decrypt_binary(const std::string& file_name,const std::string& passwd,caffe::NetParameter& params);
20 | 
21 | private:
22 |     static void encrypt_proto(const std::string& file_name,const std::string& passwd,const std::string& output_file);
23 |     static void encrypt_binary(const std::string& file_name,const std::string& passwd,const std::string& output_file);
24 | 
25 | };
26 | 
27 | #endif //GRPC_DATAFORMAT_CAFFE_ENCRYPT_H
28 | 


--------------------------------------------------------------------------------
/encrypt_tool.cpp:
--------------------------------------------------------------------------------
 1 | //
 2 | // Created by liuhao on 9/6/18.
 3 | //
 4 | 
 5 | #include "caffe_encrypt.h"
 6 | #include <string>
 7 | #include <iostream>
 8 | int main(int argc,char* argv[])
 9 | {
10 |     if(argc<4)
11 |     {
12 |         std::cerr<<"argument not enough"<<std::endl;
13 |         return -1;
14 |     }
15 | 
16 |     std::string prefix=argv[1];
17 |     std::string pwd=argv[2];
18 |     std::string output_name=argv[3];
19 | 
20 |     caffe_encrypt::encrypt(prefix,pwd,output_name);
21 | 
22 |     return 0;
23 | }


--------------------------------------------------------------------------------