Python db_session.add函数代码示例

    def add_view(self, view_name, analyzer_name, query_string=None,
                 query_dsl=None, query_filter=None):
        """Add saved view to the Sketch.

        Args:
            view_name: The name of the view.
            analyzer_name: The name of the analyzer.
            query_string: Elasticsearch query string.
            query_dsl: Dictionary with Elasticsearch DSL query.
            query_filter: Dictionary with Elasticsearch filters.

        Raises:
            ValueError: If both query_string an query_dsl are missing.

        Returns: An instance of a SQLAlchemy View object.
        """
        if not query_string or query_dsl:
            raise ValueError('Both query_string and query_dsl are missing.')

        if not query_filter:
            query_filter = {'indices': '_all'}

        name = '[{0:s}] {1:s}'.format(analyzer_name, view_name)
        view = View.get_or_create(name=name, sketch=self.sql_sketch, user=None)
        view.query_string = query_string
        view.query_filter = view.validate_filter(query_filter)
        view.query_dsl = query_dsl
        view.searchtemplate = None

        db_session.add(view)
        db_session.commit()
        return view

def explore(sketch_id, view_id=None):
    """Generates the sketch explore view template.

    Returns:
        Template with context.
    """
    sketch = Sketch.query.get_with_acl(sketch_id)
    if view_id:
        view = View.query.get(view_id)
    else:
        view = View.query.filter(
            View.user == current_user,
            View.name == u'',
            View.sketch_id == sketch_id).order_by(
                View.created_at.desc()).first()
    if not view:
        view = View(
            user=current_user, name=u'', sketch=sketch, query_string=u'',
            query_filter=u'{}')
        db_session.add(view)
        db_session.commit()
    sketch_timelines = u','.join(
        [t.searchindex.index_name for t in sketch.timelines])
    view_form = SaveViewForm()

    return render_template(
        u'sketch/explore.html', sketch=sketch, view=view,
        timelines=sketch_timelines, view_form=view_form)

    def run_wrapper(self):
        """A wrapper method to run the analyzer.

        This method is decorated to flush the bulk insert operation on the
        datastore. This makes sure that all events are indexed at exit.

        Returns:
            Return value of the run method.
        """
        result = self.run()

        # Update the searchindex description with analyzer result.
        # TODO: Don't overload the description field.
        searchindex = SearchIndex.query.filter_by(
            index_name=self.index_name).first()

        # Some code paths set the description equals to the name. Remove that
        # here to get a clean description with only analyzer results.
        if searchindex.description == searchindex.name:
            searchindex.description = ''

        # Append the analyzer result.
        if result:
            searchindex.description = '{0:s}\n{1:s}'.format(
                searchindex.description, result)
        db_session.add(searchindex)
        db_session.commit()

        return result

def _set_timeline_status(index_name, status, error_msg=None):
    """Helper function to set status for searchindex and all related timelines.

    Args:
        index_name: Name of the datastore index.
        status: Status to set.
        error_msg: Error message.
    """
    searchindex = SearchIndex.query.filter_by(index_name=index_name).first()
    timelines = Timeline.query.filter_by(searchindex=searchindex).all()

    # Set status
    searchindex.set_status(status)
    for timeline in timelines:
        timeline.set_status(status)
        db_session.add(timeline)

    # Update description if there was a failure in ingestion
    if error_msg and status == 'fail':
        # TODO: Don't overload the description field.
        searchindex.description = error_msg

    # Commit changes to database
    db_session.add(searchindex)
    db_session.commit()

def home():
    """Generates the home page view template.

    Returns:
        Template with context.
    """
    form = HiddenNameDescriptionForm()
    sketches = Sketch.all_with_acl().filter(
        not_(Sketch.Status.status == 'deleted'),
        Sketch.Status.parent).order_by(Sketch.updated_at.desc())
    # Only render upload button if it is configured.
    upload_enabled = current_app.config['UPLOAD_ENABLED']

    # Handle form for creating a new sketch.
    if form.validate_on_submit():
        sketch = Sketch(
            name=form.name.data,
            description=form.description.data,
            user=current_user)
        sketch.status.append(sketch.Status(user=None, status='new'))
        db_session.add(sketch)
        db_session.commit()

        # Give the requesting user permissions on the new sketch.
        sketch.grant_permission(permission='read', user=current_user)
        sketch.grant_permission(permission='write', user=current_user)
        sketch.grant_permission(permission='delete', user=current_user)
        return redirect(url_for('sketch_views.overview', sketch_id=sketch.id))

    return render_template(
        'home/home.html',
        sketches=sketches,
        form=form,
        upload_enabled=upload_enabled)

def timeline(sketch_id, timeline_id):
    """Generates the sketch timeline view template.

    Returns:
        Template with context.
    """
    timeline_form = TimelineForm()
    sketch = Sketch.query.get_with_acl(sketch_id)
    sketch_timeline = Timeline.query.filter(
        Timeline.id == timeline_id, Timeline.sketch == sketch).first()
    if not sketch_timeline:
        abort(HTTP_STATUS_CODE_NOT_FOUND)

    if timeline_form.validate_on_submit():
        if not sketch.has_permission(current_user, u'write'):
            abort(HTTP_STATUS_CODE_FORBIDDEN)
        sketch_timeline.name = timeline_form.name.data
        sketch_timeline.description = timeline_form.description.data
        sketch_timeline.color = timeline_form.color.data
        db_session.add(sketch_timeline)
        db_session.commit()
        return redirect(
            url_for(u'sketch_views.timeline', sketch_id=sketch.id,
                    timeline_id=sketch_timeline.id))

    return render_template(
        u'sketch/timeline.html', sketch=sketch, timeline=sketch_timeline,
        timeline_form=timeline_form)

  def WriteHeader(self):
    """Sets up the Elasticsearch index and the Timesketch database object.

    Creates the Elasticsearch index with Timesketch specific settings and the
    Timesketch SearchIndex database object.
    """
    # This cannot be static because we use the value of self._document_type
    # from arguments.
    mappings = {
        self._document_type: {
            'properties': {
                'timesketch_label': {
                    'type': 'nested'
                }
            }
        }
    }

    # Get Elasticsearch host and port from Timesketch configuration.
    with self._timesketch.app_context():
      self._host = current_app.config['ELASTIC_HOST']
      self._port = current_app.config['ELASTIC_PORT']

    self._Connect()

    self._CreateIndexIfNotExists(self._index_name, mappings)

    user = None
    if self._timeline_owner:
      user = timesketch_user.User.query.filter_by(
          username=self._timeline_owner).first()
      if not user:
        raise RuntimeError(
            'Unknown Timesketch user: {0:s}'.format(self._timeline_owner))
    else:
      logger.warning('Timeline will be visible to all Timesketch users')

    with self._timesketch.app_context():
      search_index = timesketch_sketch.SearchIndex.get_or_create(
          name=self._timeline_name, description=self._timeline_name, user=user,
          index_name=self._index_name)

      # Grant the user read permission on the mapping object and set status.
      # If user is None the timeline becomes visible to all users.
      search_index.grant_permission(user=user, permission='read')

      # In case we have a user grant additional permissions.
      if user:
        search_index.grant_permission(user=user, permission='write')
        search_index.grant_permission(user=user, permission='delete')

      # Let the Timesketch UI know that the timeline is processing.
      search_index.set_status('processing')

      # Save the mapping object to the Timesketch database.
      timesketch_db_session.add(search_index)
      timesketch_db_session.commit()

    logger.debug('Adding events to Timesketch.')

 def run(self, name):
     """Creates the group."""
     if not isinstance(name, six.text_type):
         name = codecs.decode(name, 'utf-8')
     group = Group.get_or_create(name=name)
     db_session.add(group)
     db_session.commit()
     sys.stdout.write('Group {0:s} created\n'.format(name))

    def _commit_to_database(self, model):
        """Add object to the database session and commit.

        Args:
            model: Instance of timesketch.models.[model] object
        """
        db_session.add(model)
        db_session.commit()

    def post(self, sketch_id):
        """Handles POST request to the resource.

        Args:
            sketch_id: Integer primary key for a sketch database model

        Returns:
            An annotation in JSON (instance of flask.wrappers.Response)
        """
        form = EventAnnotationForm.build(request)
        if form.validate_on_submit():
            sketch = Sketch.query.get_with_acl(sketch_id)
            indices = [t.searchindex.index_name for t in sketch.timelines]
            annotation_type = form.annotation_type.data
            searchindex_id = form.searchindex_id.data
            searchindex = SearchIndex.query.get(searchindex_id)
            event_id = form.event_id.data

            if searchindex_id not in indices:
                abort(HTTP_STATUS_CODE_BAD_REQUEST)

            def _set_label(label, toggle=False):
                """Set label on the event in the datastore."""
                self.datastore.set_label(
                    searchindex_id, event_id, sketch.id, current_user.id, label,
                    toggle=toggle)

            # Get or create an event in the SQL database to have something to
            # attach the annotation to.
            event = Event.get_or_create(
                sketch=sketch, searchindex=searchindex,
                document_id=event_id)

            # Add the annotation to the event object.
            if u'comment' in annotation_type:
                annotation = Event.Comment(
                    comment=form.annotation.data, user=current_user)
                event.comments.append(annotation)
                _set_label(u'__ts_comment')
            elif u'label' in annotation_type:
                annotation = Event.Label.get_or_create(
                    label=form.annotation.data, user=current_user)
                if annotation not in event.labels:
                    event.labels.append(annotation)
                toggle = False
                if u'__ts_star' in form.annotation.data:
                    toggle = True
                _set_label(form.annotation.data, toggle)
            else:
                abort(HTTP_STATUS_CODE_BAD_REQUEST)

            # Save the event to the database
            db_session.add(event)
            db_session.commit()

            return self.to_json(
                annotation, status_code=HTTP_STATUS_CODE_CREATED)
        return abort(HTTP_STATUS_CODE_BAD_REQUEST)

  def WriteHeader(self):
    """Setup the Elasticsearch index and the Timesketch database object.

    Creates the Elasticsearch index with Timesketch specific settings and the
    Timesketch SearchIndex database object.
    """
    # This cannot be static because we use the value of self._doc_type from
    # arguments.
    _document_mapping = {
        self._doc_type: {
            u'properties': {
                u'timesketch_label': {
                    u'type': u'nested'
                }
            }
        }
    }

    # Get Elasticsearch host and port from Timesketch configuration.
    with self._timesketch.app_context():
      _host = current_app.config[u'ELASTIC_HOST']
      _port = current_app.config[u'ELASTIC_PORT']

    self._elastic = ElasticSearchHelper(
        self._output_mediator, _host, _port, self._flush_interval,
        self._index_name, _document_mapping, self._doc_type)

    user = None
    if self._username:
      user = User.query.filter_by(username=self._username).first()
      if not user:
        raise RuntimeError(
            u'Unknown Timesketch user: {0:s}'.format(self._username))
    else:
      logging.warning(u'Timeline will be visible to all Timesketch users')

    with self._timesketch.app_context():
      search_index = SearchIndex.get_or_create(
          name=self._timeline_name, description=self._timeline_name, user=user,
          index_name=self._index_name)

      # Grant the user read permission on the mapping object and set status.
      # If user is None the timeline becomes visible to all users.
      search_index.grant_permission(user=user, permission=u'read')

      # In case we have a user grant additional permissions.
      if user:
        search_index.grant_permission(user=user, permission=u'write')
        search_index.grant_permission(user=user, permission=u'delete')

      # Let the Timesketch UI know that the timeline is processing.
      search_index.set_status(u'processing')

      # Save the mapping object to the Timesketch database.
      db_session.add(search_index)
      db_session.commit()

    logging.info(u'Adding events to Timesketch.')

  def Close(self):
    """Closes the connection to TimeSketch Elasticsearch database.

    Sends the remaining events for indexing and removes the processing status on
    the Timesketch search index object.
    """
    self._elastic.AddEvent(None, force_flush=True)
    with self._timesketch.app_context():
      search_index = SearchIndex.query.filter_by(
          index_name=self._index_name).first()
      search_index.status.remove(search_index.status[0])
      db_session.add(search_index)
      db_session.commit()

  def Close(self):
    """Closes the connection to TimeSketch Elasticsearch database.

    Sends the remaining events for indexing and removes the processing status on
    the Timesketch search index object.
    """
    super(TimesketchOutputModule, self).Close()

    with self._timesketch.app_context():
      search_index = timesketch_sketch.SearchIndex.query.filter_by(
          index_name=self._index_name).first()
      search_index.status.remove(search_index.status[0])
      timesketch_db_session.add(search_index)
      timesketch_db_session.commit()

def home():
    """Generates the home page view template.

    Returns:
        Template with context.
    """
    form = HiddenNameDescriptionForm()
    sketches = Sketch.all_with_acl().filter(
        not_(Sketch.Status.status == u'deleted'),
        Sketch.Status.parent).order_by(Sketch.updated_at.desc())
    query_filter = request.args.get(u'filter', u'')
    query = request.args.get(u'q', u'')
    # Only render upload button if it is configured.
    upload_enabled = current_app.config[u'UPLOAD_ENABLED']
    last_sketch = View.query.filter_by(
        user=current_user, name=u'').order_by(
            View.updated_at.desc()).first()

    if query_filter:
        if query_filter == u'user':
            sketches = sketches.filter(Sketch.user == current_user)
        elif query_filter == u'shared':
            sketches = sketches.filter(not_(Sketch.user == current_user))

    # TODO: Figure out a better way to handle this.
    if query:
        if query.startswith(u'*'):
            query = u''
        else:
            sketches = sketches.filter(Sketch.name.contains(query)).limit(100)

    # Handle form for creating a new sketch.
    if form.validate_on_submit():
        sketch = Sketch(
            name=form.name.data, description=form.description.data,
            user=current_user)
        sketch.status.append(sketch.Status(user=None, status=u'new'))
        # Give the requesting user permissions on the new sketch.
        sketch.grant_permission(current_user, u'read')
        sketch.grant_permission(current_user, u'write')
        sketch.grant_permission(current_user, u'delete')
        db_session.add(sketch)
        db_session.commit()
        return redirect(url_for(u'sketch_views.overview', sketch_id=sketch.id))

    return render_template(
        u'home/home.html', sketches=sketches, form=form, query=query,
        upload_enabled=upload_enabled, last_sketch=last_sketch)

    def post(self):
        """Handles POST request to the resource.

        Returns:
            A view in JSON (instance of flask.wrappers.Response)

        Raises:
            ApiHTTPError
        """
        UPLOAD_ENABLED = current_app.config[u'UPLOAD_ENABLED']
        UPLOAD_FOLDER = current_app.config[u'UPLOAD_FOLDER']

        form = UploadFileForm()
        if form.validate_on_submit() and UPLOAD_ENABLED:
            from timesketch.lib.tasks import run_plaso
            file_storage = form.file.data
            timeline_name = form.name.data
            # We do not need a human readable filename or
            # datastore index name, so we use UUIDs here.
            filename = unicode(uuid.uuid4().hex)
            index_name = unicode(uuid.uuid4().hex)

            file_path = os.path.join(UPLOAD_FOLDER, filename)
            file_storage.save(file_path)

            search_index = SearchIndex.get_or_create(
                name=timeline_name, description=timeline_name,
                user=current_user, index_name=index_name)
            search_index.grant_permission(permission=u'read', user=current_user)
            search_index.grant_permission(
                permission=u'write', user=current_user)
            search_index.grant_permission(
                permission=u'delete', user=current_user)
            search_index.set_status(u'processing')
            db_session.add(search_index)
            db_session.commit()

            run_plaso.apply_async(
                (file_path, timeline_name, index_name), task_id=index_name)

            return self.to_json(
                search_index, status_code=HTTP_STATUS_CODE_CREATED)
        else:
            raise ApiHTTPError(
                message=form.errors[u'file'][0],
                status_code=HTTP_STATUS_CODE_BAD_REQUEST)