[Forensics-changes] [yara] 14/160: Improve streams

Hilko Bengen bengen at moszumanska.debian.org
Sat Jul 1 10:29:12 UTC 2017


This is an automated email from the git hooks/post-receive script.

bengen pushed a commit to annotated tag v3.4.0
in repository yara.

commit ef7da2c52a407f22371a8744689dcc7eae4bbc7e
Author: Victor M. Alvarez <plusvic at gmail.com>
Date:   Fri Feb 13 12:44:35 2015 +0100

    Improve streams
    
    - Implement file functions in top of stream functions
    - Add support for stream writing
---
 libyara/arena.c               | 246 +++++++++---------------------------
 libyara/include/yara/arena.h  |  15 +--
 libyara/include/yara/rules.h  |   5 +
 libyara/include/yara/stream.h |  20 ++-
 libyara/libyara.sym           |   1 +
 libyara/rules.c               |  96 +++++++-------
 libyara/stream.c              |  21 +++-
 yara-python/tests.py          |  18 ++-
 yara-python/yara-python.c     | 286 +++++++++++++++++++++++++-----------------
 9 files changed, 340 insertions(+), 368 deletions(-)

diff --git a/libyara/arena.c b/libyara/arena.c
index e39946b..47fd9bd 100644
--- a/libyara/arena.c
+++ b/libyara/arena.c
@@ -873,117 +873,23 @@ int yr_arena_duplicate(
 
 
 //
-// yr_arena_save
-//
-// Saves the arena into a file. If the file exists its overwritten. This
-// function requires the arena to be coalesced.
-//
-// Args:
-//    YR_ARENA* arena          - Pointer to the arena.
-//    const char* filename     - File path.
-//
-// Returns:
-//    ERROR_SUCCESS if succeed or the corresponding error code otherwise.
-//
-
-int yr_arena_save(
-  YR_ARENA* arena,
-  const char* filename)
-{
-  YR_ARENA_PAGE* page;
-  YR_RELOC* reloc;
-  FILE* fh;
-  ARENA_FILE_HEADER header;
-
-  int32_t end_marker = -1;
-  uint8_t** reloc_address;
-  uint8_t* reloc_target;
-
-  // Only coalesced arenas can be saved.
-  assert(arena->flags & ARENA_FLAGS_COALESCED);
-
-  fh = fopen(filename, "wb");
-
-  if (fh == NULL)
-    return ERROR_COULD_NOT_OPEN_FILE;
-
-  page = arena->page_list_head;
-  reloc = page->reloc_list_head;
-
-  // Convert pointers to offsets before saving.
-  while (reloc != NULL)
-  {
-    reloc_address = (uint8_t**) (page->address + reloc->offset);
-    reloc_target = *reloc_address;
-
-    if (reloc_target != NULL)
-    {
-      assert(reloc_target >= page->address);
-      assert(reloc_target < page->address + page->used);
-      *reloc_address = (uint8_t*) (*reloc_address - page->address);
-    }
-    else
-    {
-      *reloc_address = (uint8_t*) (size_t) 0xFFFABADA;
-    }
-
-    reloc = reloc->next;
-  }
-
-  header.magic[0] = 'Y';
-  header.magic[1] = 'A';
-  header.magic[2] = 'R';
-  header.magic[3] = 'A';
-  header.size = page->size;
-  header.version = ARENA_FILE_VERSION;
-
-  fwrite(&header, sizeof(header), 1, fh);
-  fwrite(page->address, sizeof(uint8_t), header.size, fh);
-
-  reloc = page->reloc_list_head;
-
-  // Convert offsets back to pointers.
-  while (reloc != NULL)
-  {
-    fwrite(&reloc->offset, sizeof(reloc->offset), 1, fh);
-
-    reloc_address = (uint8_t**) (page->address + reloc->offset);
-    reloc_target = *reloc_address;
-
-    if (reloc_target != (void*) (size_t) 0xFFFABADA)
-      *reloc_address += (size_t) page->address;
-    else
-      *reloc_address = 0;
-
-    reloc = reloc->next;
-  }
-
-  fwrite(&end_marker, sizeof(end_marker), 1, fh);
-  fclose(fh);
-
-  return ERROR_SUCCESS;
-}
-
-
-//
-// yr_arena_load
+// yr_arena_load_stream
 //
-// Loads an arena from a file.
+// Loads an arena from a stream.
 //
 // Args:
-//    const char* filename  - File path.
-//    YR_ARENA**            - Address where a pointer to the loaded arena
-//                            will be returned.
+//    YR_STREAM* stream  - Pointer to stream object
+//    YR_ARENA**         - Address where a pointer to the loaded arena
+//                         will be returned
 //
 // Returns:
-//    ERROR_SUCCESS if succeed or the corresponding error code otherwise.
+//    ERROR_SUCCESS if successful, appropriate error code otherwise.
 //
 
-int yr_arena_load(
-    const char* filename,
+int yr_arena_load_stream(
+    YR_STREAM* stream,
     YR_ARENA** arena)
 {
-  FILE* fh;
   YR_ARENA_PAGE* page;
   YR_ARENA* new_arena;
   ARENA_FILE_HEADER header;
@@ -992,67 +898,39 @@ int yr_arena_load(
   uint8_t** reloc_address;
   uint8_t* reloc_target;
 
-  long file_size;
   int result;
 
-  fh = fopen(filename, "rb");
-
-  if (fh == NULL)
-    return ERROR_COULD_NOT_OPEN_FILE;
-
-  fseek(fh, 0, SEEK_END);
-  file_size = ftell(fh);
-  fseek(fh, 0, SEEK_SET);
-
-  if (fread(&header, sizeof(header), 1, fh) != 1)
-  {
-    fclose(fh);
+  if (yr_stream_read(&header, sizeof(header), 1, stream) != 1)
     return ERROR_INVALID_FILE;
-  }
 
   if (header.magic[0] != 'Y' ||
       header.magic[1] != 'A' ||
       header.magic[2] != 'R' ||
       header.magic[3] != 'A')
   {
-    fclose(fh);
     return ERROR_INVALID_FILE;
   }
 
-  if (header.size >= file_size)
-  {
-    fclose(fh);
-    return ERROR_CORRUPT_FILE;
-  }
-
   if (header.version > ARENA_FILE_VERSION)
-  {
-    fclose(fh);
     return ERROR_UNSUPPORTED_FILE_VERSION;
-  }
 
   result = yr_arena_create(header.size, 0, &new_arena);
 
   if (result != ERROR_SUCCESS)
-  {
-    fclose(fh);
     return result;
-  }
 
   page = new_arena->current_page;
 
-  if (fread(page->address, header.size, 1, fh) != 1)
+  if (yr_stream_read(page->address, header.size, 1, stream) != 1)
   {
-    fclose(fh);
     yr_arena_destroy(new_arena);
     return ERROR_CORRUPT_FILE;
   }
 
   page->used = header.size;
 
-  if (fread(&reloc_offset, sizeof(reloc_offset), 1, fh) != 1)
+  if (yr_stream_read(&reloc_offset, sizeof(reloc_offset), 1, stream) != 1)
   {
-    fclose(fh);
     yr_arena_destroy(new_arena);
     return ERROR_CORRUPT_FILE;
   }
@@ -1069,101 +947,101 @@ int yr_arena_load(
     else
       *reloc_address = 0;
 
-    if (fread(&reloc_offset, sizeof(reloc_offset), 1, fh) != 1)
+    if (yr_stream_read(&reloc_offset, sizeof(reloc_offset), 1, stream) != 1)
     {
-      fclose(fh);
       yr_arena_destroy(new_arena);
       return ERROR_CORRUPT_FILE;
     }
   }
 
-  fclose(fh);
-
   *arena = new_arena;
 
   return ERROR_SUCCESS;
 }
 
+
 //
-// yr_arena_load_stream
+// yr_arena_save_stream
 //
-// Loads an arena from a stream.
+// Saves the arena into a stream. If the file exists its overwritten. This
+// function requires the arena to be coalesced.
 //
 // Args:
-//    YR_STREAM* stream  - stream object
-//    YR_ARENA**         - Address where a pointer to the loaded arena
-//                         will be returned
+//    YR_ARENA* arena         - Pointer to the arena.
+//    YR_STREAM* stream       - Pointer to stream object.
 //
 // Returns:
-//    ERROR_SUCCESS if successful, appropriate error code otherwise.
+//    ERROR_SUCCESS if succeed or the corresponding error code otherwise.
 //
 
-int yr_arena_load_stream(
-    YR_STREAM* stream,
-    YR_ARENA** arena)
+int yr_arena_save_stream(
+  YR_ARENA* arena,
+  YR_STREAM* stream)
 {
   YR_ARENA_PAGE* page;
-  YR_ARENA* new_arena;
+  YR_RELOC* reloc;
   ARENA_FILE_HEADER header;
 
-  int32_t reloc_offset;
+  int32_t end_marker = -1;
   uint8_t** reloc_address;
   uint8_t* reloc_target;
 
-  int result;
-
-  if (yr_stream_read(&header, sizeof(header), 1, stream) != 1)
-    return ERROR_INVALID_FILE;
-
-  if (header.magic[0] != 'Y' ||
-      header.magic[1] != 'A' ||
-      header.magic[2] != 'R' ||
-      header.magic[3] != 'A')
-    return ERROR_INVALID_FILE;
-
-  if (header.version > ARENA_FILE_VERSION)
-    return ERROR_UNSUPPORTED_FILE_VERSION;
+  // Only coalesced arenas can be saved.
+  assert(arena->flags & ARENA_FLAGS_COALESCED);
 
-  result = yr_arena_create(header.size, 0, &new_arena);
+  page = arena->page_list_head;
+  reloc = page->reloc_list_head;
 
-  if (result != ERROR_SUCCESS)
-    return result;
+  // Convert pointers to offsets before saving.
+  while (reloc != NULL)
+  {
+    reloc_address = (uint8_t**) (page->address + reloc->offset);
+    reloc_target = *reloc_address;
 
-  page = new_arena->current_page;
+    if (reloc_target != NULL)
+    {
+      assert(reloc_target >= page->address);
+      assert(reloc_target < page->address + page->used);
+      *reloc_address = (uint8_t*) (*reloc_address - page->address);
+    }
+    else
+    {
+      *reloc_address = (uint8_t*) (size_t) 0xFFFABADA;
+    }
 
-  if (yr_stream_read(page->address, header.size, 1, stream) != 1) {
-    yr_arena_destroy(new_arena);
-    return ERROR_CORRUPT_FILE;
+    reloc = reloc->next;
   }
 
-  page->used = header.size;
+  header.magic[0] = 'Y';
+  header.magic[1] = 'A';
+  header.magic[2] = 'R';
+  header.magic[3] = 'A';
+  header.size = page->size;
+  header.version = ARENA_FILE_VERSION;
 
-  if (yr_stream_read(&reloc_offset, sizeof(reloc_offset), 1, stream) != 1)
-  {
-    yr_arena_destroy(new_arena);
-    return ERROR_CORRUPT_FILE;
-  }
+  yr_stream_write(&header, sizeof(header), 1, stream);
+  yr_stream_write(page->address, sizeof(uint8_t), header.size, stream);
 
-  while (reloc_offset != -1)
+  reloc = page->reloc_list_head;
+
+  // Convert offsets back to pointers.
+  while (reloc != NULL)
   {
-    yr_arena_make_relocatable(new_arena, page->address, reloc_offset, EOL);
+    yr_stream_write(&reloc->offset, sizeof(reloc->offset), 1, stream);
 
-    reloc_address = (uint8_t**) (page->address + reloc_offset);
+    reloc_address = (uint8_t**) (page->address + reloc->offset);
     reloc_target = *reloc_address;
 
-    if (reloc_target != (uint8_t*) (size_t) 0xFFFABADA)
+    if (reloc_target != (void*) (size_t) 0xFFFABADA)
       *reloc_address += (size_t) page->address;
     else
       *reloc_address = 0;
 
-    if (yr_stream_read(&reloc_offset, sizeof(reloc_offset), 1, stream) != 1)
-    {
-      yr_arena_destroy(new_arena);
-      return ERROR_CORRUPT_FILE;
-    }
+    reloc = reloc->next;
   }
 
-  *arena = new_arena;
+  yr_stream_write(&end_marker, sizeof(end_marker), 1, stream);
 
   return ERROR_SUCCESS;
 }
+
diff --git a/libyara/include/yara/arena.h b/libyara/include/yara/arena.h
index dc1e329..de2cee3 100644
--- a/libyara/include/yara/arena.h
+++ b/libyara/include/yara/arena.h
@@ -131,21 +131,16 @@ int yr_arena_append(
     YR_ARENA* source_arena);
 
 
-int yr_arena_save(
-    YR_ARENA* arena,
-    const char* filename);
-
-
-int yr_arena_load(
-    const char* filename,
-    YR_ARENA** arena);
-
-
 int yr_arena_load_stream(
     YR_STREAM* stream,
     YR_ARENA** arena);
 
 
+int yr_arena_save_stream(
+  YR_ARENA* arena,
+  YR_STREAM* stream);
+
+
 int yr_arena_duplicate(
     YR_ARENA* arena,
     YR_ARENA** duplicated);
diff --git a/libyara/include/yara/rules.h b/libyara/include/yara/rules.h
index b8b814c..8a091fe 100644
--- a/libyara/include/yara/rules.h
+++ b/libyara/include/yara/rules.h
@@ -88,6 +88,11 @@ YR_API int yr_rules_save(
     const char* filename);
 
 
+YR_API int yr_rules_save_stream(
+    YR_RULES* rules,
+    YR_STREAM* stream);
+
+
 YR_API int yr_rules_load(
     const char* filename,
     YR_RULES** rules);
diff --git a/libyara/include/yara/stream.h b/libyara/include/yara/stream.h
index d54a9d3..2802696 100644
--- a/libyara/include/yara/stream.h
+++ b/libyara/include/yara/stream.h
@@ -22,22 +22,38 @@ limitations under the License.
 typedef size_t (*YR_STREAM_READ_FUNC)(
     void* ptr,
     size_t size,
-    size_t nmemb,
+    size_t count,
+    void* user_data);
+
+
+typedef size_t (*YR_STREAM_WRITE_FUNC)(
+    const void* ptr,
+    size_t size,
+    size_t count,
     void* user_data);
 
 
 typedef struct _YR_STREAM
 {
   void* user_data;
+
   YR_STREAM_READ_FUNC read;
+  YR_STREAM_WRITE_FUNC write;
+
 } YR_STREAM;
 
 
 size_t yr_stream_read(
     void* ptr,
     size_t size,
-    size_t nmemb,
+    size_t count,
     YR_STREAM* stream);
 
 
+size_t yr_stream_write(
+    const void* ptr,
+    size_t size,
+    size_t count,
+    YR_STREAM* stream);
+
 #endif
diff --git a/libyara/libyara.sym b/libyara/libyara.sym
index fde7cb9..3182cfc 100644
--- a/libyara/libyara.sym
+++ b/libyara/libyara.sym
@@ -22,6 +22,7 @@ yr_rules_scan_mem
 yr_rules_scan_file
 yr_rules_scan_proc
 yr_rules_save
+yr_rules_save_stream
 yr_rules_load
 yr_rules_load_stream
 yr_rules_destroy
diff --git a/libyara/rules.c b/libyara/rules.c
index 0791bdd..a1a7156 100644
--- a/libyara/rules.c
+++ b/libyara/rules.c
@@ -594,30 +594,16 @@ YR_API int yr_rules_scan_proc(
 }
 
 
-YR_API int yr_rules_save(
-    YR_RULES* rules,
-    const char* filename)
-{
-  assert(rules->tidx_mask == 0);
-  return yr_arena_save(rules->arena, filename);
-}
-
-
-YR_API int yr_rules_load(
-  const char* filename,
-  YR_RULES** rules)
+YR_API int yr_rules_load_stream(
+    YR_STREAM* stream,
+    YR_RULES** rules)
 {
-  YR_RULES* new_rules;
-  YARA_RULES_FILE_HEADER* header;
-
-  int result;
-
-  new_rules = (YR_RULES*) yr_malloc(sizeof(YR_RULES));
+  YR_RULES* new_rules = (YR_RULES*) yr_malloc(sizeof(YR_RULES));
 
   if (new_rules == NULL)
     return ERROR_INSUFICIENT_MEMORY;
 
-  result = yr_arena_load(filename, &new_rules->arena);
+  int result = yr_arena_load_stream(stream, &new_rules->arena);
 
   if (result != ERROR_SUCCESS)
   {
@@ -625,7 +611,9 @@ YR_API int yr_rules_load(
     return result;
   }
 
-  header = (YARA_RULES_FILE_HEADER*) yr_arena_base_address(new_rules->arena);
+  YARA_RULES_FILE_HEADER* header = (YARA_RULES_FILE_HEADER*)
+      yr_arena_base_address(new_rules->arena);
+
   new_rules->automaton = header->automaton;
   new_rules->code_start = header->code_start;
   new_rules->externals_list_head = header->externals_list_head;
@@ -650,50 +638,54 @@ YR_API int yr_rules_load(
 }
 
 
-YR_API int yr_rules_load_stream(
-  YR_STREAM* stream,
-  YR_RULES** rules)
+YR_API int yr_rules_load(
+    const char* filename,
+    YR_RULES** rules)
 {
-  YR_RULES* new_rules;
-  YARA_RULES_FILE_HEADER* header;
+  FILE* fh = fopen(filename, "rb");
 
-  int result;
+  if (fh == NULL)
+    return ERROR_COULD_NOT_OPEN_FILE;
 
-  new_rules = (YR_RULES*) yr_malloc(sizeof(YR_RULES));
+  YR_STREAM stream = {
+    .user_data = fh,
+    .read = (YR_STREAM_READ_FUNC) fread
+  };
 
-  if (new_rules == NULL)
-    return ERROR_INSUFICIENT_MEMORY;
+  int result = yr_rules_load_stream(&stream, rules);
 
-  result = yr_arena_load_stream(stream, &new_rules->arena);
+  fclose(fh);
+  return result;
+}
 
-  if (result != ERROR_SUCCESS)
-  {
-    yr_free(new_rules);
-    return result;
-  }
 
-  header = (YARA_RULES_FILE_HEADER*) yr_arena_base_address(new_rules->arena);
-  new_rules->automaton = header->automaton;
-  new_rules->code_start = header->code_start;
-  new_rules->externals_list_head = header->externals_list_head;
-  new_rules->rules_list_head = header->rules_list_head;
-  new_rules->tidx_mask = 0;
+YR_API int yr_rules_save_stream(
+    YR_RULES* rules,
+    YR_STREAM* stream)
+{
+  assert(rules->tidx_mask == 0);
+  return yr_arena_save_stream(rules->arena, stream);
+}
 
-  #if _WIN32
-  new_rules->mutex = CreateMutex(NULL, FALSE, NULL);
 
-  if (new_rules->mutex == NULL)
-    return ERROR_INTERNAL_FATAL_ERROR;
-  #else
-  result = pthread_mutex_init(&new_rules->mutex, NULL);
+YR_API int yr_rules_save(
+    YR_RULES* rules,
+    const char* filename)
+{
+  FILE* fh = fopen(filename, "wb");
 
-  if (result != 0)
-    return ERROR_INTERNAL_FATAL_ERROR;
-  #endif
+  if (fh == NULL)
+    return ERROR_COULD_NOT_OPEN_FILE;
 
-  *rules = new_rules;
+  YR_STREAM stream = {
+    .user_data = fh,
+    .write = (YR_STREAM_WRITE_FUNC) fwrite,
+  };
 
-  return ERROR_SUCCESS;
+  int result = yr_rules_save_stream(rules, &stream);
+
+  fclose(fh);
+  return result;
 }
 
 
diff --git a/libyara/stream.c b/libyara/stream.c
index 61dff99..a39e385 100644
--- a/libyara/stream.c
+++ b/libyara/stream.c
@@ -17,11 +17,28 @@ limitations under the License.
 #include <stddef.h>
 #include <yara/stream.h>
 
+
 size_t yr_stream_read(
     void* ptr,
     size_t size,
-    size_t nmemb,
+    size_t count,
     YR_STREAM* stream)
 {
-  return stream->read(ptr, size, nmemb, stream->user_data);
+  if (stream->read == NULL)
+  	return 0;
+
+  return stream->read(ptr, size, count, stream->user_data);
 }
+
+
+size_t yr_stream_write(
+    const void* ptr,
+    size_t size,
+    size_t count,
+    YR_STREAM* stream)
+{
+  if (stream->write == NULL)
+  	return 0;
+
+  return stream->write(ptr, size, count, stream->user_data);
+}
\ No newline at end of file
diff --git a/yara-python/tests.py b/yara-python/tests.py
index 60edf6a..e0fb777 100644
--- a/yara-python/tests.py
+++ b/yara-python/tests.py
@@ -926,13 +926,25 @@ class TestYara(unittest.TestCase):
 
         # Python 2/3
         try:
-            yac = StringIO.StringIO(YAC_FILE)
+            yac1 = StringIO.StringIO(YAC_FILE)
+            yac2 = StringIO.StringIO()
         except:
-            yac = io.BytesIO(YAC_FILE)
+            yac1 = io.BytesIO(YAC_FILE)
+            yac2 = io.BytesIO()
+
+        r = yara.load(yac1)
+        r.save(yac2)
 
-        r = yara.load(yac)
         m = r.match(data="dummy")
         self.assertTrue(len(m) == 1)
 
+        yac2.seek(0)
+        r = yara.load(yac2)
+
+        m = r.match(data="dummy")
+        self.assertTrue(len(m) == 1)
+
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/yara-python/yara-python.c b/yara-python/yara-python.c
index 0d1a718..e045518 100644
--- a/yara-python/yara-python.c
+++ b/yara-python/yara-python.c
@@ -65,7 +65,6 @@ For complete documentation please visit:\n\
 https://plusvic.github.io/yara\n"
 
 
-
 // Match object
 
 typedef struct
@@ -597,6 +596,74 @@ int yara_callback(
 }
 
 
+/* YR_STREAM read method for "file-like objects" */
+
+static size_t flo_read(
+    void* ptr,
+    size_t size,
+    size_t count,
+    void* user_data)
+{
+  for (int i = 0; i < count; i++)
+  {
+    PyGILState_STATE gil_state = PyGILState_Ensure();
+
+    PyObject* bytes = PyObject_CallMethod(
+        (PyObject*) user_data, "read", "n", (Py_ssize_t) size);
+
+    PyGILState_Release(gil_state);
+
+    if (bytes != NULL)
+    {
+      Py_ssize_t len;
+      char* buffer;
+
+      int result = PyBytes_AsStringAndSize(bytes, &buffer, &len);
+
+      Py_DECREF(bytes);
+
+      if (result == -1 || len < size)
+        return i;
+
+      memcpy(ptr + i * size, buffer, size);
+    }
+    else
+    {
+      return i;
+    }
+  }
+
+  return count;
+}
+
+
+/* YR_STREAM write method for "file-like objects" */
+
+static size_t flo_write(
+    const void* ptr,
+    size_t size,
+    size_t count,
+    void* user_data)
+{
+  for (int i = 0; i < count; i++)
+  {
+    PyGILState_STATE gil_state = PyGILState_Ensure();
+
+    PyObject* result = PyObject_CallMethod(
+        (PyObject*) user_data, "write", "s#", ptr + i * size, size);
+
+    PyGILState_Release(gil_state);
+
+    if (result == NULL)
+      return i;
+
+    Py_DECREF(result);
+  }
+
+  return count;
+}
+
+
 int process_compile_externals(
     PyObject* externals,
     YR_COMPILER* compiler)
@@ -929,6 +996,7 @@ static PyObject * Rules_next(PyObject *self)
   rule = PyObject_NEW(Rule, &Rule_Type);
   tag_list = PyList_New(0);
   meta_list = PyDict_New();
+
   if (rule != NULL && tag_list != NULL && meta_list != NULL)
   {
     yr_rule_tags_foreach(rules->iter_current_rule, tag)
@@ -1137,27 +1205,51 @@ static PyObject * Rules_save(
     PyObject *self,
     PyObject *args)
 {
-  char* filepath;
   int error;
+
+  PyObject* param;
   Rules* rules = (Rules*) self;
 
-  if (PyArg_ParseTuple(args, "s", &filepath))
+  if (!PyArg_UnpackTuple(args, "save", 1, 1, &param))
+  {
+    return PyErr_Format(
+        PyExc_TypeError,
+          "save() takes 1 argument");
+  }
+
+  if (PY_STRING_CHECK(param))
   {
+    char* filepath = PY_STRING_TO_C(param);
+
     Py_BEGIN_ALLOW_THREADS
     error = yr_rules_save(rules->rules, filepath);
     Py_END_ALLOW_THREADS
 
     if (error != ERROR_SUCCESS)
       return handle_error(error, filepath);
+  }
+  else if (PyObject_HasAttrString(param, "write"))
+  {
+    YR_STREAM stream = {
+      .user_data = param,
+      .write = flo_write
+    };
 
-    Py_RETURN_NONE;
+    Py_BEGIN_ALLOW_THREADS;
+    error = yr_rules_save_stream(rules->rules, &stream);
+    Py_END_ALLOW_THREADS;
+
+    if (error != ERROR_SUCCESS)
+      return handle_error(error, "<file-like-object>");
   }
   else
   {
     return PyErr_Format(
-        PyExc_TypeError,
-          "save() takes 1 argument");
+      PyExc_TypeError,
+      "load() expects either a file path or a file-like object");
   }
+
+  Py_RETURN_NONE;
 }
 
 
@@ -1540,135 +1632,99 @@ static PyObject * yara_compile(
 }
 
 
-/* YR_STREAM read method for "file-like objects" */
-static size_t flo_read(
-  void* ptr,
-  size_t size,
-  size_t nmemb,
-  void* user_data)
-{
-  PyObject* obj = (PyObject*)user_data;
-  int i;
-  for (i = 0; i < nmemb; i++)
-  {
-    PyGILState_STATE gil_state = PyGILState_Ensure();
-    PyObject* rv = PyObject_CallMethod(obj, "read", "n", (Py_ssize_t)size);
-    PyGILState_Release(gil_state);
-    if (rv == NULL)
-      return i;
-    Py_ssize_t len;
-    char* buf;
-    if (PyBytes_AsStringAndSize(rv, &buf, &len) == -1)
-      return i;
-    if (len < size)
-      return i;
-    memcpy(ptr + i*size, buf, size);
-  }
-  return nmemb;
-}
-
-
 static PyObject * yara_load(
     PyObject *self,
     PyObject *args)
 {
-  YR_EXTERNAL_VARIABLE* external;
+  Rules* rules = PyObject_NEW(Rules, &Rules_Type);
+  PyObject* param;
 
-  int error;
+  if (rules == NULL)
+    return PyErr_NoMemory();
 
-  Rules* rules;
+  if (!PyArg_UnpackTuple(args, "load", 1, 1, &param))
+  {
+    return PyErr_Format(
+        PyExc_TypeError,
+          "load() takes 1 argument");
+  }
 
-  PyObject* param;
+  int error;
 
-  if (PyArg_UnpackTuple(args, "load", 1, 1, &param))
+  if (PY_STRING_CHECK(param))
   {
-    rules = PyObject_NEW(Rules, &Rules_Type);
-    if (rules == NULL)
-      return PyErr_NoMemory();
+    char* filepath = PY_STRING_TO_C(param);
 
-    if (PY_STRING_CHECK(param))
-    {
-      char* filepath = PY_STRING_TO_C(param);
-      Py_BEGIN_ALLOW_THREADS;
-      error = yr_rules_load(filepath, &rules->rules);
-      Py_END_ALLOW_THREADS;
+    Py_BEGIN_ALLOW_THREADS;
+    error = yr_rules_load(filepath, &rules->rules);
+    Py_END_ALLOW_THREADS;
 
-      if (error != ERROR_SUCCESS)
-        return handle_error(error, filepath);
-    }
-    else
-    {
-      if (PyObject_HasAttrString(param, "read"))
-      {
-        YR_STREAM stream = {
-          .user_data = param,
-          .read = flo_read
-        };
+    if (error != ERROR_SUCCESS)
+      return handle_error(error, filepath);
+  }
+  else if (PyObject_HasAttrString(param, "read"))
+  {
+    YR_STREAM stream = {
+      .user_data = param,
+      .read = flo_read
+    };
 
-        Py_BEGIN_ALLOW_THREADS;
-        error = yr_rules_load_stream(&stream, &rules->rules);
-        Py_END_ALLOW_THREADS;
+    Py_BEGIN_ALLOW_THREADS;
+    error = yr_rules_load_stream(&stream, &rules->rules);
+    Py_END_ALLOW_THREADS;
 
-        if (error != ERROR_SUCCESS)
-          return handle_error(error, "<file-like-object>");
-      }
-      else
-      {
-        return PyErr_Format(
-          PyExc_TypeError,
-          "can't use argument");
-      }
-    }
+    if (error != ERROR_SUCCESS)
+      return handle_error(error, "<file-like-object>");
+  }
+  else
+  {
+    return PyErr_Format(
+      PyExc_TypeError,
+      "load() expects either a file path or a file-like object");
+  }
 
-    external = rules->rules->externals_list_head;
-    rules->iter_current_rule = rules->rules->rules_list_head;
+  YR_EXTERNAL_VARIABLE* external = rules->rules->externals_list_head;
+  rules->iter_current_rule = rules->rules->rules_list_head;
 
-    if (!EXTERNAL_VARIABLE_IS_NULL(external))
-      rules->externals = PyDict_New();
-    else
-      rules->externals = NULL;
+  if (!EXTERNAL_VARIABLE_IS_NULL(external))
+    rules->externals = PyDict_New();
+  else
+    rules->externals = NULL;
 
-    while (!EXTERNAL_VARIABLE_IS_NULL(external))
+  while (!EXTERNAL_VARIABLE_IS_NULL(external))
+  {
+    switch(external->type)
     {
-      switch(external->type)
-      {
-        case EXTERNAL_VARIABLE_TYPE_BOOLEAN:
-          PyDict_SetItemString(
-              rules->externals,
-              external->identifier,
-              PyBool_FromLong((long) external->value.i));
-          break;
-        case EXTERNAL_VARIABLE_TYPE_INTEGER:
-          PyDict_SetItemString(
-              rules->externals,
-              external->identifier,
-              PyLong_FromLong((long) external->value.i));
-          break;
-        case EXTERNAL_VARIABLE_TYPE_FLOAT:
-          PyDict_SetItemString(
-              rules->externals,
-              external->identifier,
-              PyFloat_FromDouble(external->value.f));
-          break;
-        case EXTERNAL_VARIABLE_TYPE_STRING:
-          PyDict_SetItemString(
-              rules->externals,
-              external->identifier,
-              PY_STRING(external->value.s));
-          break;
-      }
-
-      external++;
+      case EXTERNAL_VARIABLE_TYPE_BOOLEAN:
+        PyDict_SetItemString(
+            rules->externals,
+            external->identifier,
+            PyBool_FromLong((long) external->value.i));
+        break;
+      case EXTERNAL_VARIABLE_TYPE_INTEGER:
+        PyDict_SetItemString(
+            rules->externals,
+            external->identifier,
+            PyLong_FromLong((long) external->value.i));
+        break;
+      case EXTERNAL_VARIABLE_TYPE_FLOAT:
+        PyDict_SetItemString(
+            rules->externals,
+            external->identifier,
+            PyFloat_FromDouble(external->value.f));
+        break;
+      case EXTERNAL_VARIABLE_TYPE_STRING:
+        PyDict_SetItemString(
+            rules->externals,
+            external->identifier,
+            PY_STRING(external->value.s));
+        break;
     }
 
-    return (PyObject*) rules;
-  }
-  else
-  {
-    return PyErr_Format(
-        PyExc_TypeError,
-          "load() takes 1 argument");
+    external++;
   }
+
+  return (PyObject*) rules;
 }
 
 

-- 
Alioth's /usr/local/bin/git-commit-notice on /srv/git.debian.org/git/forensics/yara.git



More information about the forensics-changes mailing list