Repository: kevoreilly/CAPEv2

Issue:
can't reschedule tasks once temporary files are gone
# Prerequisites

Please answer the following questions for yourself before submitting an issue.

- [x] I am running the latest version
- [x] I checked the documentation and found no answer
- [x] I checked to make sure that this issue has not already been filed
- [x] I'm reporting the issue to the correct repository (for multi-repository projects)
- [x] I'm have read all configs with all optional parts


# Expected Behavior

Rescheduling a task succeeds after its temporary file has been removed.

# Current Behavior

Rescheduling a task currently uses [task.target](https://github.com/kevoreilly/CAPEv2/blob/fb4283b540b3293595670b49e1576811960c8f59/lib/cuckoo/core/database.py#L1838), which is a temporary file that could be gone removed at any point. When that happens the task can no longer be rescheduled.

# Failure Information (for bugs)

## Steps to Reproduce

1. submit a task
1. delete its temporary file
1. reschedule it

## Request Failures

```json
{"error": true, "error_value": "An error occured while trying to reschedule Task ID 1234"}
```


Discussion:

Resolving patch:
diff --git a/lib/cuckoo/core/database.py b/lib/cuckoo/core/database.py
--- a/lib/cuckoo/core/database.py
+++ b/lib/cuckoo/core/database.py
@@ -1834,8 +1834,30 @@ def reschedule(self, task_id):
         else:
             tags = task.tags
 
+        def _ensure_valid_target(task):
+            if task.category == "url":
+                # URL tasks always have valid targets, return it as-is.
+                return task.target
+
+            # All other task types have a "target" pointing to a temp location,
+            # so get a stable path "target" based on the sample hash.
+            paths = self.sample_path_by_hash(task.sample.sha256)
+            paths = [x for x in paths if os.path.exists(x)]
+            if not paths:
+                return None
+
+            if task.category == 'pcap':
+                # PCAP task paths are represented as bytes
+                return paths[0].encode()
+            return paths[0]
+
+        task_target = _ensure_valid_target(task)
+        if not task_target:
+            log.warning("Unable to find valid target for task: %s", task_id)
+            return
+
         return add(
-            task.target,
+            task_target,
             task.timeout,
             task.package,
             task.options,


Test patch:
diff --git a/tests/test_database.py b/tests/test_database.py
--- a/tests/test_database.py
+++ b/tests/test_database.py
@@ -3,10 +3,13 @@
 # See the file 'docs/LICENSE' for copying permission.
 
 from __future__ import absolute_import
+import base64
 import os
+import shutil
 from tempfile import NamedTemporaryFile
 
 from lib.cuckoo.core.database import Database, Task
+from lib.cuckoo.common.utils import store_temp_file
 
 
 class TestDatabaseEngine(object):
@@ -17,14 +20,20 @@ class TestDatabaseEngine(object):
     def setup_method(self):
         with NamedTemporaryFile(mode="w+", delete=False) as f:
             f.write("hehe")
-        self.temp = f.name
+        self.temp_filename = f.name
+        pcap_header_base64 = b"1MOyoQIABAAAAAAAAAAAAAAABAABAAAA"
+        pcap_bytes = base64.b64decode(pcap_header_base64)
+        self.temp_pcap = store_temp_file(pcap_bytes, "%s.pcap" % f.name)
         self.d = Database(dsn="sqlite://")
         # self.d.connect(dsn=self.URI)
         self.session = self.d.Session()
+        self.binary_storage = os.path.join(os.getcwd(), "storage/binaries")
+        os.makedirs(self.binary_storage)
 
     def teardown_method(self):
         del self.d
-        os.unlink(self.temp)
+        os.unlink(self.temp_filename)
+        shutil.rmtree(self.binary_storage)
 
     def add_url(self, url, priority=1, status="pending"):
         task_id = self.d.add_url(url, priority=priority)
@@ -35,7 +44,7 @@ def test_add_tasks(self):
 
         # Add task.
         count = self.session.query(Task).count()
-        self.d.add_path(self.temp)
+        self.d.add_path(self.temp_filename)
         assert self.session.query(Task).count() == count + 1
 
         # Add url.
@@ -56,15 +65,82 @@ def test_long_error(self):
         assert err and len(err[0].message) == 1024
 
     def test_task_set_options(self):
-        assert self.d.add_path(self.temp, options={"foo": "bar"}) is None
-        t1 = self.d.add_path(self.temp, options="foo=bar")
+        assert self.d.add_path(self.temp_filename, options={"foo": "bar"}) is None
+        t1 = self.d.add_path(self.temp_filename, options="foo=bar")
         assert self.d.view_task(t1).options == "foo=bar"
 
     def test_task_tags_str(self):
-        task = self.d.add_path(self.temp, tags="foo,,bar")
+        task = self.d.add_path(self.temp_filename, tags="foo,,bar")
         tag_list = list(self.d.view_task(task).tags)
         assert [str(x.name) for x in tag_list].sort() == ["foo", "bar"].sort()
 
+    def test_reschedule_file(self):
+        count = self.session.query(Task).count()
+        task_id = self.d.add_path(self.temp_filename)
+        assert self.session.query(Task).count() == count + 1
+        task = self.d.view_task(task_id)
+        assert task is not None
+        assert task.category == "file"
+
+        # write a real sample to storage
+        sample_path = os.path.join(self.binary_storage, task.sample.sha256)
+        shutil.copy(self.temp_filename, sample_path)
+
+        new_task_id = self.d.reschedule(task_id)
+        assert new_task_id is not None
+        new_task = self.d.view_task(new_task_id)
+        assert new_task.category == "file"
+
+    def test_reschedule_static(self):
+        count = self.session.query(Task).count()
+        task_id = self.d.add_static(self.temp_filename)
+        assert self.session.query(Task).count() == count + 1
+        task = self.d.view_task(task_id)
+        assert task is not None
+        assert task.category == "static"
+
+        # write a real sample to storage
+        static_path = os.path.join(self.binary_storage, task.sample.sha256)
+        shutil.copy(self.temp_filename, static_path)
+
+        new_task_id = self.d.reschedule(task_id)
+        assert new_task_id is not None
+        new_task = self.d.view_task(new_task_id)
+        assert new_task.category == "static"
+
+    def test_reschedule_pcap(self):
+        count = self.session.query(Task).count()
+        task_id = self.d.add_pcap(self.temp_pcap)
+        assert self.session.query(Task).count() == count + 1
+        task = self.d.view_task(task_id)
+        assert task is not None
+        assert task.category == "pcap"
+
+        # write a real sample to storage
+        pcap_path = os.path.join(self.binary_storage, task.sample.sha256)
+        shutil.copy(self.temp_pcap, pcap_path)
+
+        # reschedule the PCAP task
+        new_task_id = self.d.reschedule(task_id)
+        assert new_task_id is not None
+        new_task = self.d.view_task(new_task_id)
+        assert new_task.category == "pcap"
+
+    def test_reschedule_url(self):
+        # add a URL task
+        count = self.session.query(Task).count()
+        task_id = self.d.add_url("test_reschedule_url")
+        assert self.session.query(Task).count() == count + 1
+        task = self.d.view_task(task_id)
+        assert task is not None
+        assert task.category == "url"
+
+        # reschedule the URL task
+        new_task_id = self.d.reschedule(task_id)
+        assert new_task_id is not None
+        new_task = self.d.view_task(new_task_id)
+        assert new_task.category == "url"
+
     def test_add_machine(self):
         self.d.add_machine(
             name="name1",