Exploit heuristics & task limit

owl-cli/src/exploit.rs

@@ -16,6 +16,14 @@ fn arg_match_numeric_none(matches: &ArgMatches, name: &str) -> Result<Option<i32
     }
 }
 
+fn arg_match_numeric_disable(matches: &ArgMatches, name: &str) -> Result<i32, Error> {
+    if matches.is_present(name) && matches.value_of(name).unwrap() != "disable" {
+        Ok(matches.value_of(name).unwrap().parse::<i32>()?)
+    } else {
+        Ok(-1)
+    }
+}
+
 pub fn exploit_command() -> App<'static, 'static> {
     SubCommand::with_name("exploit")
         .about("User defined exploit management")
@@ -87,6 +95,13 @@ pub fn exploit_command() -> App<'static, 'static> {
                 .args(&[
                     Arg::from_usage("<name> 'name of the exploit to download'"),
                 ]),
+            SubCommand::with_name("failure")
+                .about("update failure value for specific exploit target (admin)")
+                .args(&[
+                    Arg::from_usage("<exploit_name> 'name of the exploit to modify'"),
+                    Arg::from_usage("<service_variant_name> 'name of the targeting service variant'"),
+                    Arg::from_usage("<failure> 'failure value to set'"),
+                ]),
             SubCommand::with_name("run-all")
                 .about("trigger scheduled exploit process (admin)"),
             SubCommand::with_name("run")
@@ -282,6 +297,22 @@ pub fn exploit_match(matches: &ArgMatches, shared_param: SharedParam) -> Result<
             Ok("Exploit successfully downloaded".to_string())
         },
 
+        ("failure", Some(matches)) => {
+            shared_param.client.failure_exploit(
+                shared_param.token,
+                ExploitFailureParams {
+                    exploit_name: matches.value_of("exploit_name").unwrap().to_string(),
+                    service_variant_name: matches
+                        .value_of("service_variant_name")
+                        .unwrap()
+                        .to_string(),
+                    failure: arg_match_numeric_disable(matches, "failure")?,
+                },
+            )?;
+
+            Ok("Exploit failure successfully updated".to_string())
+        },
+
         ("run-all", Some(_matches)) => {
             shared_param.client.run_all_exploit(shared_param.token)?;
 

owl-daemon/Cargo.toml

@@ -20,6 +20,7 @@ shell-escape = "0.1.4"
 tarpc = { git = "https://github.com/PLUS-POSTECH/tarpc.git" }
 tokio = "0.1.7"
 tokio-core = "0.1.17"
+tokio-threadpool = "0.1.5"
 toml = "0.4"
 r2d2 = "0.8.2"
 r2d2-diesel = "1.0.0"

owl-daemon/migrations/2018-08-11-222334_service_target_heuristics/down.sql

@@ -0,0 +1 @@
+ALTER TABLE exploit_targets DROP COLUMN consecutive_failure;

owl-daemon/migrations/2018-08-11-222334_service_target_heuristics/up.sql

@@ -0,0 +1 @@
+ALTER TABLE exploit_targets ADD COLUMN consecutive_failure integer;

owl-daemon/src/db/models.rs

@@ -168,6 +168,21 @@ pub struct ExploitAttachmentInsertable {
 pub struct ExploitTarget {
     pub exploit_id: i32,
     pub service_variant_id: i32,
+    pub consecutive_failure: i32,
+}
+
+#[derive(Insertable)]
+#[table_name = "exploit_targets"]
+pub struct ExploitTargetInsertable {
+    pub exploit_id: i32,
+    pub service_variant_id: i32,
+    pub consecutive_failure: i32,
+}
+
+#[derive(AsChangeset)]
+#[table_name = "exploit_targets"]
+pub struct ExploitTargetChangeset {
+    pub consecutive_failure: Option<i32>,
 }
 
 #[derive(DbEnum, Debug)]

owl-daemon/src/db/schema.rs

@@ -26,6 +26,7 @@ table! {
     exploit_targets (exploit_id, service_variant_id) {
         exploit_id -> Int4,
         service_variant_id -> Int4,
+        consecutive_failure -> Int4,
     }
 }
 

owl-daemon/src/exploit.rs

@@ -127,6 +127,7 @@ impl ExploitManager {
     pub fn exploit_future<F, T: AsRef<Path>>(
         &self,
         db_pool: DbPool,
+        exploit_task_id: i32,
         exploit_name: T,
         entry_file: &str,
         target_info: TargetInfo,
@@ -164,6 +165,10 @@ impl ExploitManager {
             Deadline::new(
                 update(db_pool.clone(), (ExploitStatus::Running, String::new()))
                 .into_future()
+                .and_then(move |_| {
+                    info!(target: "exploit", "[ExploitTask] Started: {}", exploit_task_id);
+                    Ok(())
+                })
                 // run exploit
                 .and_then(move |_| shell_command_runner(&full_command)
                     .current_dir(exploit_directory)

owl-daemon/src/handler/exploit.rs

@@ -3,7 +3,9 @@ use db;
 use db::models::*;
 use db::schema::*;
 use diesel;
+use diesel::dsl::exists;
 use diesel::prelude::*;
+use diesel::select;
 use diesel::PgConnection;
 use digest::Input;
 use error::Error;
@@ -201,6 +203,49 @@ pub fn edit_exploit(resource: &DaemonResource, params: ExploitEditParams) -> Res
     }
 }
 
+pub fn failure_exploit(
+    resource: &DaemonResource,
+    params: ExploitFailureParams,
+) -> Result<(), Error> {
+    let con: &PgConnection = &*resource.db_pool.get()?;
+    let exploit_name = params.exploit_name;
+    let service_variant_name = params.service_variant_name;
+    let failure = params.failure;
+
+    con.transaction::<(), Error, _>(|| {
+        let exploit = exploits::table
+            .filter(exploits::name.eq(&exploit_name))
+            .first::<Exploit>(con)?;
+
+        let service_variant = service_variants::table
+            .filter(service_variants::name.eq(&service_variant_name))
+            .first::<ServiceVariant>(con)?;
+
+        let query = exploit_targets::table
+            .filter(exploit_targets::exploit_id.eq(exploit.id))
+            .filter(exploit_targets::service_variant_id.eq(service_variant.service_id));
+
+        if select(exists(query)).get_result(con)? {
+            diesel::update(query)
+                .set(ExploitTargetChangeset {
+                    consecutive_failure: Some(failure),
+                })
+                .execute(con)?;
+            info!(target: "db", "[ExploitTarget] Update record: ({}, {}) -> {}", &exploit.name, &service_variant.name, failure);
+        } else {
+            diesel::insert_into(exploit_targets::table)
+                .values(ExploitTargetInsertable {
+                    exploit_id: exploit.id,
+                    service_variant_id: service_variant.id,
+                    consecutive_failure: failure,
+                })
+                .execute(con)?;
+            info!(target: "db", "[ExploitTarget] Insert record: ({}, {}) -> {}", &exploit.name, &service_variant.name, failure);
+        }
+        Ok(())
+    })
+}
+
 pub fn download_exploit(
     resource: &DaemonResource,
     params: ExploitDownloadParams,
@@ -245,11 +290,22 @@ fn create_exploit_task(
         .get_result::<ExploitTask>(con)?)
 }
 
-fn active_target(con: &PgConnection, service_id: i32) -> Result<Vec<(TargetInfo, i32)>, Error> {
+fn active_target(
+    con: &PgConnection,
+    failure_threshold: i32,
+    exploit: &Exploit,
+) -> Result<Vec<(TargetInfo, i32)>, Error> {
+    let target_blacklist = exploit_targets::table
+        .filter(exploit_targets::exploit_id.eq(exploit.id))
+        .filter(exploit_targets::consecutive_failure.gt(failure_threshold))
+        .select(exploit_targets::service_variant_id)
+        .load::<i32>(con)?;
+
     Ok(service_providers::table
         .inner_join(teams::table)
         .inner_join(service_variants::table.inner_join(services::table))
-        .filter(services::id.eq(service_id))
+        .filter(services::id.eq(exploit.service_id))
+        .filter(service_variants::id.ne_all(target_blacklist))
         .order_by((
             service_variants::service_id,
             service_providers::team_id,
@@ -289,48 +345,111 @@ fn update_exploit_status(
             exploit_tasks::message.eq(exploit_message),
         ))
         .execute(con)?;
-    info!(target: "exploit", "[ExploitTask] Status updated: {} - {:?}", &exploit_task_id, &exploit_status);
+    info!(target: "exploit", "[ExploitTask] Status updated: {} - {:?}", exploit_task_id, &exploit_status);
     Ok(result)
 }
 
+fn update_exploit_consecutive_failure(
+    con: &PgConnection,
+    exploit_task_id: i32,
+    failed: bool,
+) -> Result<(), Error> {
+    con.transaction::<(), Error, _>(|| {
+        let (exploit_task, service_provider) = exploit_tasks::table
+            .filter(exploit_tasks::id.eq(exploit_task_id))
+            .inner_join(service_providers::table)
+            .first::<(ExploitTask, ServiceProvider)>(con)?;
+
+        let exploit_target = exploit_targets::table
+            .find((
+                &exploit_task.exploit_id,
+                &service_provider.service_variant_id,
+            ))
+            .first::<ExploitTarget>(con);
+
+        match exploit_target {
+            Ok(exploit_target) => {
+                let failure_count = exploit_target.consecutive_failure;
+                let new_failure_count = failure_count + 1;
+                let query = diesel::update(
+                    exploit_targets::table
+                        .filter(exploit_targets::exploit_id.eq(exploit_task.exploit_id))
+                        .filter(
+                            exploit_targets::service_variant_id
+                                .eq(service_provider.service_variant_id),
+                        ),
+                );
+
+                if failed {
+                    if failure_count < 0 {
+                        info!(target: "exploit", "[ExploitTarget] Failure not counted: ({}, {})", exploit_task.exploit_id, service_provider.service_variant_id);
+                        Ok(())
+                    } else {
+                        info!(target: "exploit", "[ExploitTarget] Failure count incremented: ({}, {}) - {} times", exploit_task.exploit_id, service_provider.service_variant_id, new_failure_count);
+                        query
+                            .set(exploit_targets::consecutive_failure.eq(new_failure_count))
+                            .execute(con)?;
+                        Ok(())
+                    }
+                } else {
+                    info!(target: "exploit", "[ExploitTarget] Failure count reset: {}", exploit_task.exploit_id);
+                    query
+                        .set(exploit_targets::consecutive_failure.eq(0))
+                        .execute(con)?;
+                    Ok(())
+                }
+            },
+            Err(_) => {
+                if failed {
+                    info!(target: "exploit", "[ExploitTarget] Failure count incremented: {} - {} times", exploit_task.exploit_id, 1);
+                    diesel::insert_into(exploit_targets::table)
+                        .values(ExploitTargetInsertable {
+                            exploit_id: exploit_task.exploit_id,
+                            service_variant_id: service_provider.service_variant_id,
+                            consecutive_failure: 1,
+                        })
+                        .execute(con)?;
+                    Ok(())
+                } else {
+                    Ok(())
+                }
+            },
+        }
+    })
+}
+
 fn add_to_task_queue(resource: &DaemonResource, exploit: &Exploit) -> Result<(), Error> {
     let con: &PgConnection = &*resource.db_pool.get()?;
+    let failure_threshold = resource.config.exploit_config.failure_threshold;
 
-    // TODO: exclude failing exploits
-    let _exploit_task_targets = exploit_targets::table
-        .inner_join(service_variants::table)
-        .filter(exploit_targets::exploit_id.eq(exploit.id))
-        .filter(service_variants::service_id.eq(exploit.service_id))
-        .load::<(ExploitTarget, ServiceVariant)>(con)?;
-
-    // TODO: only load exploit name
-    let exploit_attachments = exploit_attachments::table
+    let exploit_attachment_name = exploit_attachments::table
         .filter(exploit_attachments::exploit_id.eq(exploit.id))
-        .load::<ExploitAttachment>(con)?;
+        .order(exploit_attachments::id.asc())
+        .select(exploit_attachments::name)
+        .first::<String>(con)?;
 
     let timeout = exploit
         .timeout
         .unwrap_or(resource.config.exploit_config.default_timeout);
 
-    for (target_info, service_provider_id) in active_target(con, exploit.service_id)? {
+    for (target_info, service_provider_id) in active_target(con, failure_threshold, exploit)? {
         let exploit_task = create_exploit_task(
             con,
             exploit,
             service_provider_id,
             resource.config.exploit_config.default_retries,
         )?;
 
-        info!(target: "exploit", "[ExploitTask] Generated: {}", &exploit_task.id);
-
         let db_pool = resource.db_pool.clone();
         let exploit_task_id = exploit_task.id;
         resource.task_executor.spawn(
             resource
                 .exploit_manager
                 .exploit_future(
                     resource.db_pool.clone(),
+                    exploit.id,
                     exploit.name.clone(),
-                    &exploit_attachments[0].name,
+                    &exploit_attachment_name,
                     target_info,
                     timeout,
                     move |db_pool, exploit_update_message| {
@@ -345,8 +464,9 @@ fn add_to_task_queue(resource: &DaemonResource, exploit: &Exploit) -> Result<(),
                 .then(move |result| {
                     if let Ok(con) = db_pool.get() {
                         let con = &*con;
+                        let failed = result.is_err();
                         let exploit_update_message = result.unwrap_or_else(|err| {
-                            info!(target: "exploit", "[ExploitTask] Failed: {}", &exploit_task.id);
+                            info!(target: "exploit", "[ExploitTask] Failed: {}", exploit_task.id);
                             if let Error::ExploitError(err) = err {
                                 match err {
                                     ExploitError::ExploitProcessError(err) => {
@@ -382,12 +502,13 @@ fn add_to_task_queue(resource: &DaemonResource, exploit: &Exploit) -> Result<(),
                         });
 
                         update_exploit_status(con, exploit_task_id, exploit_update_message).is_ok();
+                        update_exploit_consecutive_failure(con, exploit_task_id, failed).is_ok();
                     }
 
                     Ok(())
                 }),
         );
-        info!(target: "exploit", "[ExploitTask] Spawned: {}", &exploit_task_id);
+        info!(target: "exploit", "[ExploitTask] Spawned: {}", exploit_task_id);
     }
 
     Ok(())

owl-daemon/src/lib.rs

@@ -60,6 +60,8 @@ pub struct ExploitConfig {
     pub auth_command: String,
     pub default_retries: i32,
     pub default_timeout: i32,
+    pub failure_threshold: i32,
+    pub max_running_exploit_task: i32,
 }
 
 #[derive(Clone)]
@@ -333,6 +335,21 @@ impl FutureService for OwlDaemon {
         )
     }
 
+    type FailureExploitFut = Result<(), Message>;
+    fn failure_exploit(
+        &self,
+        cli_token: String,
+        params: ExploitFailureParams,
+    ) -> Self::FailureExploitFut {
+        run_handler_with_param(
+            Permission::Admin,
+            cli_token,
+            handler::exploit::failure_exploit,
+            &self.resource,
+            params,
+        )
+    }
+
     type ListExploitFut = Result<Vec<ExploitData>, Message>;
     fn list_exploit(&self, cli_token: String, params: ExploitListParams) -> Self::ListExploitFut {
         run_handler_with_param(