Changes to storage interface to enforce storage types (log2timeline#2939

)

plaso/storage/interface.py

@@ -445,11 +445,16 @@ def WriteSessionCompletion(self, session_completion):
       session_completion (SessionCompletion): session completion information.
 
     Raises:
-      IOError: when the storage file is closed or read-only.
-      OSError: when the storage file is closed or read-only.
+      IOError: if the storage type does not support writing a session
+          completion or the storage file is closed or read-only.
+      OSError: if the storage type does not support writing a session
+          completion or the storage file is closed or read-only.
     """
     self._RaiseIfNotWritable()
 
+    if self.storage_type != definitions.STORAGE_TYPE_SESSION:
+      raise IOError('Session completion not supported by storage type.')
+
     self._WriteAttributeContainer(session_completion)
 
   def WriteSessionStart(self, session_start):
@@ -459,11 +464,16 @@ def WriteSessionStart(self, session_start):
       session_start (SessionStart): session start information.
 
     Raises:
-      IOError: when the storage file is closed or read-only.
-      OSError: when the storage file is closed or read-only.
+      IOError: if the storage type does not support writing a session
+          start or the storage file is closed or read-only.
+      OSError: if the storage type does not support writing a session
+          start or the storage file is closed or read-only.
     """
     self._RaiseIfNotWritable()
 
+    if self.storage_type != definitions.STORAGE_TYPE_SESSION:
+      raise IOError('Session start not supported by storage type.')
+
     self._WriteAttributeContainer(session_start)
 
   def WriteTaskCompletion(self, task_completion):
@@ -473,11 +483,16 @@ def WriteTaskCompletion(self, task_completion):
       task_completion (TaskCompletion): task completion information.
 
     Raises:
-      IOError: when the storage file is closed or read-only.
-      OSError: when the storage file is closed or read-only.
+      IOError: if the storage type does not support writing a task
+          completion or the storage file is closed or read-only.
+      OSError: if the storage type does not support writing a task
+          completion or the storage file is closed or read-only.
     """
     self._RaiseIfNotWritable()
 
+    if self.storage_type != definitions.STORAGE_TYPE_TASK:
+      raise IOError('Task start not supported by storage type.')
+
     self._WriteAttributeContainer(task_completion)
 
   def WriteTaskStart(self, task_start):
@@ -487,10 +502,16 @@ def WriteTaskStart(self, task_start):
       task_start (TaskStart): task start information.
 
     Raises:
-      StorageNotReadableError: when the storage file is closed or read-only.
+      IOError: if the storage type does not support writing a task
+          start or the storage file is closed or read-only.
+      OSError: if the storage type does not support writing a task
+          start or the storage file is closed or read-only.
     """
     self._RaiseIfNotWritable()
 
+    if self.storage_type != definitions.STORAGE_TYPE_TASK:
+      raise IOError('Task completion not supported by storage type.')
+
     self._WriteAttributeContainer(task_start)
 
   def _DeserializeAttributeContainer(self, container_type, serialized_data):

tests/storage/sqlite/sqlite_file.py

@@ -542,7 +542,7 @@ def testWriteSessionStartAndCompletion(self):
     with shared_test_lib.TempDirectory() as temp_directory:
       temp_file = os.path.join(temp_directory, 'plaso.sqlite')
       storage_file = sqlite_file.SQLiteStorageFile(
-          storage_type=definitions.STORAGE_TYPE_TASK)
+          storage_type=definitions.STORAGE_TYPE_SESSION)
       storage_file.Open(path=temp_file, read_only=False)
 
       storage_file.WriteSessionStart(session_start)