1 : // Copyright 2012 Google Inc. All Rights Reserved.
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // http://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 : //
15 : // Decomposes an image and serializes the decomposed image to file.
16 :
17 : #include "syzygy/pe/decompose_app.h"
18 :
19 : #include "base/at_exit.h"
20 : #include "base/command_line.h"
21 : #include "base/files/file_path.h"
22 : #include "base/files/file_util.h"
23 : #include "base/strings/string_util.h"
24 : #include "base/time/time.h"
25 : #include "syzygy/block_graph/block_graph.h"
26 : #include "syzygy/pe/decomposer.h"
27 : #include "syzygy/pe/pe_file.h"
28 : #include "syzygy/pe/serialization.h"
29 :
30 : namespace pe {
31 :
32 : using block_graph::BlockGraph;
33 : using block_graph::BlockGraphSerializer;
34 : using application::ScopedTimeLogger;
35 :
36 : namespace {
37 :
38 : const char kUsageFormatStr[] =
39 : "Usage: %ls [options]\n"
40 : "\n"
41 : " A tool that uses symbol information and disassembly to decompose a\n"
42 : " PE image file into discrete blocks of code (and data), and to infer\n"
43 : " the references between them, serializing the resulting decomposition\n"
44 : " for later use.\n"
45 : "\n"
46 : "Required parameters\n"
47 : " --image=<image file>\n"
48 : "Optional parameters\n"
49 : " --benchmark-load\n"
50 : " Causes the output to be deserialized after serialization,\n"
51 : " for benchmarking.\n"
52 : " --graph-only\n"
53 : " Causes the serialized output to only contain the block-graph, with\n"
54 : " all data inlined. The PE file (and pe_lib) will not be needed to\n"
55 : " deserialize the resulting file. Useful for producing canned unittest\n"
56 : " data.\n"
57 : " --output=<output file>\n"
58 : " The location of output file. If not specified, will append\n"
59 : " '.bg' to the image file.\n"
60 : " --strip-strings\n"
61 : " If specified then the serialized decomposition will not contain any\n"
62 : " strings.\n";
63 :
64 : } // namespace
65 :
66 : void DecomposeApp::PrintUsage(const base::FilePath& program,
67 E : const base::StringPiece& message) {
68 E : if (!message.empty()) {
69 E : ::fwrite(message.data(), 1, message.length(), out());
70 E : ::fprintf(out(), "\n\n");
71 : }
72 :
73 E : ::fprintf(out(), kUsageFormatStr, program.BaseName().value().c_str());
74 E : }
75 :
76 E : bool DecomposeApp::ParseCommandLine(const base::CommandLine* cmd_line) {
77 E : DCHECK(cmd_line != NULL);
78 :
79 E : if (cmd_line->HasSwitch("help")) {
80 E : PrintUsage(cmd_line->GetProgram(), "");
81 E : return false;
82 : }
83 :
84 E : image_path_ = cmd_line->GetSwitchValuePath("image");
85 E : if (image_path_.empty()) {
86 E : PrintUsage(cmd_line->GetProgram(), "Must specify '--image' parameter!");
87 E : return false;
88 : }
89 :
90 : // If no output file specified, use default.
91 E : output_path_ = cmd_line->GetSwitchValuePath("output");
92 E : if (output_path_.empty()) {
93 E : output_path_ = base::FilePath(image_path_.value() + L".bg");
94 E : LOG(INFO) << "Inferring output path from image path.";
95 : }
96 :
97 E : benchmark_load_ = cmd_line->HasSwitch("benchmark-load");
98 E : graph_only_ = cmd_line->HasSwitch("graph-only");
99 E : strip_strings_ = cmd_line->HasSwitch("strip-strings");
100 :
101 E : return true;
102 E : }
103 :
104 E : int DecomposeApp::Run() {
105 E : LOG(INFO) << "Processing \"" << image_path_.value() << "\".";
106 :
107 : // Parse the PE File.
108 E : pe::PEFile pe_file;
109 : {
110 E : ScopedTimeLogger scoped_time_logger("Parsing PE file");
111 E : if (!pe_file.Init(image_path_))
112 i : return 1;
113 E : }
114 :
115 : // Decompose the image.
116 E : BlockGraph block_graph;
117 E : pe::ImageLayout image_layout(&block_graph);
118 E : pe::Decomposer decomposer(pe_file);
119 : {
120 E : ScopedTimeLogger scoped_time_logger("Decomposing image");
121 E : if (!decomposer.Decompose(&image_layout))
122 i : return 1;
123 E : }
124 :
125 : // Save the decomposition do the output path.
126 : {
127 E : ScopedTimeLogger scoped_time_logger("Saving decomposed image");
128 E : if (!SaveDecomposedImage(pe_file, image_layout, output_path_))
129 i : return 1;
130 E : }
131 :
132 : // If requested, benchmark the time it takes to reload the decomposition.
133 E : if (benchmark_load_) {
134 E : ScopedTimeLogger scoped_time_logger("Loading decomposed image");
135 E : if (!LoadDecomposedImage(output_path_))
136 i : return 1;
137 E : }
138 :
139 E : return 0;
140 E : }
141 :
142 : bool DecomposeApp::SaveDecomposedImage(
143 : const pe::PEFile& pe_file, const pe::ImageLayout& image_layout,
144 E : const base::FilePath& output_path) const {
145 E : base::ScopedFILE out_file(base::OpenFile(output_path, "wb"));
146 E : core::FileOutStream out_stream(out_file.get());
147 E : core::NativeBinaryOutArchive out_archive(&out_stream);
148 :
149 E : BlockGraphSerializer::Attributes attributes = 0;
150 E : if (strip_strings_)
151 E : attributes |= BlockGraphSerializer::OMIT_STRINGS;
152 :
153 E : if (graph_only_) {
154 E : BlockGraphSerializer bgs;
155 E : bgs.set_attributes(attributes);
156 E : bgs.set_data_mode(BlockGraphSerializer::OUTPUT_ALL_DATA);
157 E : if (!bgs.Save(*image_layout.blocks.graph(), &out_archive)) {
158 i : LOG(ERROR) << "Unable to save block-graph.";
159 i : return false;
160 : }
161 E : } else {
162 : if (!SaveBlockGraphAndImageLayout(pe_file, attributes, image_layout,
163 E : &out_archive)) {
164 i : LOG(ERROR) << "Unable to save image decomposition.";
165 i : return false;
166 : }
167 : }
168 :
169 E : if (!out_archive.Flush())
170 i : return false;
171 :
172 E : return true;
173 E : }
174 :
175 E : bool DecomposeApp::LoadDecomposedImage(const base::FilePath& file_path) const {
176 E : pe::PEFile pe_file;
177 E : BlockGraph block_graph;
178 :
179 E : base::ScopedFILE in_file(base::OpenFile(file_path, "rb"));
180 E : core::FileInStream in_stream(in_file.get());
181 E : core::NativeBinaryInArchive in_archive(&in_stream);
182 :
183 E : if (graph_only_) {
184 E : BlockGraphSerializer bgs;
185 E : if (!bgs.Load(&block_graph, &in_archive)) {
186 i : LOG(ERROR) << "Unable to load block-graph.";
187 i : return false;
188 : }
189 E : } else {
190 E : pe::ImageLayout image_layout(&block_graph);
191 E : BlockGraphSerializer::Attributes attributes = 0;
192 : if (!LoadBlockGraphAndImageLayout(&pe_file, &attributes, &image_layout,
193 E : &in_archive)) {
194 i : LOG(ERROR) << "Unable to load image decomposition.";
195 i : return false;
196 : }
197 E : }
198 :
199 E : LOG(INFO) << "Successfully loaded image decomposition.";
200 E : return true;
201 E : }
202 :
203 : } // namespace pe
|