Adding screen option

JLospinoso · Sep 27, 2017 · add6f4b · add6f4b
1 parent a629917
commit add6f4b
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 16 deletions.
diff --git a/Action.h b/Action.h
@@ -27,8 +27,8 @@ struct HeadAction {
 };
 
 struct GetAction {
-  GetAction(const std::string& path_dir, bool is_verbose)
-    : re{"[^a-zA-Z0-9.-]"}, is_verbose{is_verbose}, path_dir {path_dir} {
+  GetAction(const std::string& path_dir, const std::string& screen, bool is_verbose)
+    : re{ "[^a-zA-Z0-9.-]" }, is_verbose{ is_verbose }, path_dir{ path_dir }, screen{ screen } {
     boost::filesystem::path boost_path(path_dir);
     boost::system::error_code ec;
     boost::filesystem::create_directories(path_dir, ec);
@@ -53,14 +53,25 @@ struct GetAction {
     auto path{path_dir};
     path.append("/");
     path.append(regex_replace(target, re, "_"));
-
-    std::ofstream file;
-    file.exceptions(std::ios_base::failbit | std::ios_base::badbit);
-    file.open(path, std::ofstream::out);
-    file << contents << std::endl;
+    if (screen.empty()) {
+      std::ofstream file;
+      file.exceptions(std::ios_base::failbit | std::ios_base::badbit);
+      file.open(path, std::ofstream::out);
+      file << contents << std::endl;
+    } else {
+      std::stringstream ss;
+      ss << contents << std::endl;
+      const auto contents_str = ss.str();
+      const auto contains_screen = contents_str.find(screen) != std::string::npos;
+      if (contains_screen) return;
+      std::ofstream file;
+      file.exceptions(std::ios_base::failbit | std::ios_base::badbit);
+      file.open(path, std::ofstream::out);
+      file << contents_str << std::endl;
+    }
   }
 
   const boost::regex re;
   const bool is_verbose;
-  const std::string path_dir;
+  const std::string path_dir, screen;
 };
diff --git a/Options.cpp b/Options.cpp
@@ -19,6 +19,7 @@ Options::Options(int argc, const char** argv)
     ("out", value<string>(&output_path)->default_value(""), "output path. dir if contents enabled. (default: HOSTNAME)")
     ("err", value<string>(&error_path)->default_value(""), "error path (file). (default: HOSTNAME-err.log)")
     ("proxy", value<string>(&proxy)->default_value(""), "SOCKS5 proxy address:port. (default: none)")
+    ("screen", value<string>(&screen)->default_value(""), "omits 200-level response if contents contains screen (default: none)")
     ("stdin,d", bool_switch(&from_stdin), "read from stdin (default: no)")
     ("tls,t", bool_switch(&tls), "use tls/ssl (default: no)")
     ("sensitive,s", bool_switch(&sensitive_teardown), "complain about rude TCP teardowns (default: no)")
@@ -88,6 +89,7 @@ string Options::get_pretty_print() const noexcept {
     "[ ] User Agent: " << get_user_agent() << "\n" <<
     "[ ] Proxy: " << (is_proxy() ? get_proxy() : "No") << "\n" <<
     "[ ] Contents: " << (is_contents() ? "Yes" : "No") << "\n" <<
+    "[ ] Screen: " << (get_screen().empty() ? "None" : get_screen()) << "\n" <<
     "[ ] Output: " << get_output_path() << "\n" <<
     "[ ] Error Output: " << get_error_path() << "\n" <<
     "[ ] Verbose: " << (is_verbose() ? "Yes" : "No") << "\n" <<
@@ -143,6 +145,8 @@ const string& Options::get_host() const noexcept { return host; }
 
 const string& Options::get_user_agent() const noexcept { return user_agent; }
 
+const string& Options::get_screen() const noexcept { return screen; }
+
 size_t Options::get_initial_coroutines() const noexcept { return initial_coroutines; }
 
 size_t Options::get_minimum_coroutines() const noexcept { return minimum_coroutines; }

diff --git a/Options.h b/Options.h
@@ -25,6 +25,7 @@ struct Options {
   const std::string& get_output_path() const noexcept;
   const std::string& get_error_path() const noexcept;
   const std::string& get_user_agent() const noexcept;
+  const std::string& get_screen() const noexcept;
   size_t get_initial_coroutines() const noexcept;
   size_t get_minimum_coroutines() const noexcept;
   size_t get_maximum_coroutines() const noexcept;
@@ -34,7 +35,7 @@ struct Options {
   size_t initial_coroutines, minimum_coroutines, maximum_coroutines, sample_size, sample_interval;
   bool help, tls, verify, contents, verbose, optimize, print_found, tor, sensitive_teardown, 
         leading_zeros, test, telescoping, from_stdin;
-  std::string host, pattern, output_path, error_path, help_str, proxy, user_agent;
+  std::string host, pattern, output_path, error_path, help_str, proxy, user_agent, screen;
 };
 
 struct OptionsException : std::runtime_error {

diff --git a/main.cpp b/main.cpp
@@ -36,7 +36,7 @@ namespace {
 
   GetQuery make_get(const Options& options) {
     return GetQuery{
-      GetAction{ options.get_output_path(), options.is_verbose() }, options.is_print_found(), options.is_verbose()
+      GetAction{ options.get_output_path(), options.get_screen(), options.is_verbose() }, options.is_print_found(), options.is_verbose()
     };
   }
 
@@ -45,11 +45,6 @@ namespace {
       HeadAction{ options.get_output_path(), options.is_verbose() }, options.is_print_found(), options.is_verbose()
     };
   }
-
-  template <typename Generator>
-  Generator make_generator(const Options& options) {
-
-  }
 }
 
 int main(int argc, const char** argv) {
@@ -73,7 +68,7 @@ int main(int argc, const char** argv) {
     if (!options.is_stdin()) {
       UriGenerator uri_generator{ options.get_pattern(), options.is_leading_zeros(), options.is_telescoping() };
       try {
-        auto cardinality = uri_generator.get_range_size();
+        const auto cardinality = uri_generator.get_range_size();
         cout << "[ ] URL generation set cardinality is " << cardinality << endl;
       }
       catch (const overflow_error&) {