API Documentation

@add_docs(
    description="""
    REST endpoint to fetch list of job comments or
    retrieve/delete a job comment with job comment ID.
    Requires authentication.
    """
)
class CommentViewSet(ModelViewSet):
    """
    CommentViewSet provides the following actions:

    - **list**: Retrieve a list of comments associated with jobs visible to the authenticated user.
    - **retrieve**: Retrieve a specific comment by ID, accessible to the comment's owner or anyone in the same organization.
    - **destroy**: Delete a comment by ID, allowed only for the comment's owner.
    - **update**: Update a comment by ID, allowed only for the comment's owner.
    - **partial_update**: Partially update a comment by ID, allowed only for the comment's owner.

    Permissions:
    - **IsAuthenticated**: Requires authentication for all actions.
    - **IsObjectUserPermission**: Allows only the comment owner to update or delete the comment.
    - **IsObjectUserOrSameOrgPermission**: Allows the comment owner or anyone in the same organization to retrieve the comment.

    Queryset:
    - Filters comments to include only those associated with jobs visible to the authenticated user.
    """

    queryset = Comment.objects.all()
    serializer_class = CommentSerializer
    permission_classes = [IsAuthenticated]

    def get_permissions(self):
        """
        Customizes permissions based on the action being performed.

        - For `destroy`, `update`, and `partial_update` actions, adds `IsObjectUserPermission` to ensure that only
          the comment owner can perform these actions.
        - For the `retrieve` action, adds `IsObjectUserOrSameOrgPermission` to allow the comment owner or anyone in the same
          organization to retrieve the comment.

        Returns:
        - List of applicable permissions.
        """
        permissions = super().get_permissions()

        # only the owner of the comment can update or delete the comment
        if self.action in ["destroy", "update", "partial_update"]:
            permissions.append(IsObjectUserPermission())
        # the owner and anyone in the org can read the comment
        if self.action in ["retrieve"]:
            permissions.append(IsObjectUserOrSameOrgPermission())

        return permissions

    def get_queryset(self):
        """
        Filters the queryset to include only comments related to jobs visible to the authenticated user.

        - Fetches job IDs that are visible to the user.
        - Filters the comment queryset to include only comments associated with these jobs.

        Returns:
        - Filtered queryset of comments.
        """
        queryset = super().get_queryset()
        jobs = Job.objects.visible_for_user(self.request.user).values_list(
            "pk", flat=True
        )
        return queryset.filter(job__id__in=jobs)

@add_docs(
    description="""
    REST endpoint to fetch list of jobs or retrieve/delete a job with job ID.
    Requires authentication.
    """
)
class JobViewSet(ReadAndDeleteOnlyViewSet, SerializerActionMixin):
    """
    JobViewSet provides the following actions:

    - **list**: Retrieve a list of jobs visible to the authenticated user, ordered by request time.
    - **retrieve**: Retrieve a specific job by ID.
    - **destroy**: Delete a job by ID, allowed only for the job owner or anyone in the same organization.
    - **recent_scans**: Retrieve recent jobs based on an MD5 hash, limited by a maximum temporal distance.
    - **recent_scans_user**: Retrieve recent jobs for the authenticated user, filtered by sample status.
    - **retry**: Retry a job if its status is in a final state.
    - **kill**: Kill a running job by closing celery tasks and marking it as killed.
    - **download_sample**: Download a file/sample associated with a job.
    - **pivot**: Perform a pivot operation from a job's reports.
    - **aggregate_status**: Aggregate jobs by their status over a specified time range.
    - **aggregate_type**: Aggregate jobs by type (file or observable) over a specified time range.
    - **aggregate_observable_classification**: Aggregate jobs by observable classification over a specified time range.
    - **aggregate_file_mimetype**: Aggregate jobs by file MIME type over a specified time range.
    - **aggregate_observable_name**: Aggregate jobs by observable name over a specified time range.
    - **aggregate_md5**: Aggregate jobs by MD5 hash over a specified time range.

    Permissions:
    - **IsAuthenticated**: Requires authentication for all actions.
    - **IsObjectUserOrSameOrgPermission**: Allows job deletion or killing only by the job owner or anyone in the same organization.

    Queryset:
    - Prefetches related tags and orders jobs by request time, filtered to include only jobs visible to the authenticated user.
    """

    queryset = (
        Job.objects.prefetch_related("tags").order_by("-received_request_time").all()
    )
    serializer_class = RestJobSerializer
    serializer_action_classes = {
        "retrieve": RestJobSerializer,
        "list": JobListSerializer,
    }
    filterset_class = JobFilter
    ordering_fields = [
        "received_request_time",
        "finished_analysis_time",
        "process_time",
    ]

    def get_permissions(self):
        """
        Customizes permissions based on the action being performed.

        - For `destroy` and `kill` actions, adds `IsObjectUserOrSameOrgPermission` to ensure that only
          the job owner or anyone in the same organization can perform these actions.

        Returns:
        - List of applicable permissions.
        """
        permissions = super().get_permissions()
        if self.action in ["destroy", "kill"]:
            permissions.append(IsObjectUserOrSameOrgPermission())
        return permissions

    def get_queryset(self):
        """
        Filters the queryset to include only jobs visible to the authenticated user, ordered by request time.

        Logs the request parameters and returns the filtered queryset.

        Returns:
        - Filtered queryset of jobs.
        """
        user = self.request.user
        logger.info(
            f"user: {user} request the jobs with params: {self.request.query_params}"
        )
        return Job.objects.visible_for_user(user).order_by("-received_request_time")

    @action(detail=False, methods=["post"])
    def recent_scans(self, request):
        """
        Retrieve recent jobs based on an MD5 hash, filtered by a maximum temporal distance.

        Expects the following parameters in the request data:
        - `md5`: The MD5 hash to filter jobs by.
        - `max_temporal_distance`: The maximum number of days to look back for recent jobs (default is 14 days).

        Returns:
        - List of recent jobs matching the MD5 hash.
        """
        if "md5" not in request.data:
            raise ValidationError({"detail": "md5 is required"})
        max_temporal_distance = request.data.get("max_temporal_distance", 14)
        jobs = (
            Job.objects.filter(md5=request.data["md5"])
            .visible_for_user(self.request.user)
            .filter(
                finished_analysis_time__gte=now()
                - datetime.timedelta(days=max_temporal_distance)
            )
            .annotate_importance(request.user)
            .order_by("-importance", "-finished_analysis_time")
        )
        return Response(
            JobRecentScanSerializer(jobs, many=True).data, status=status.HTTP_200_OK
        )

    @action(detail=False, methods=["post"])
    def recent_scans_user(self, request):
        """
        Retrieve recent jobs for the authenticated user, filtered by sample status.

        Expects the following parameters in the request data:
        - `is_sample`: Whether to filter jobs by sample status (required).
        - `limit`: The maximum number of recent jobs to return (default is 5).

        Returns:
        - List of recent jobs for the user.
        """
        limit = request.data.get("limit", 5)
        if "is_sample" not in request.data:
            raise ValidationError({"detail": "is_sample is required"})
        jobs = (
            Job.objects.filter(user__pk=request.user.pk)
            .filter(is_sample=request.data["is_sample"])
            .annotate_importance(request.user)
            .order_by("-importance", "-finished_analysis_time")[:limit]
        )
        return Response(
            JobRecentScanSerializer(jobs, many=True).data, status=status.HTTP_200_OK
        )

    @action(detail=True, methods=["patch"])
    def retry(self, request, pk=None):
        """
        Retry a job if its status is in a final state.

        If the job is currently running, raises a validation error.

        Returns:
        - No content (204) if the job is successfully retried.
        """
        job = self.get_object()
        if job.status not in Job.Status.final_statuses():
            raise ValidationError({"detail": "Job is running"})
        job.retry()
        return Response(status=status.HTTP_204_NO_CONTENT)

    @add_docs(
        description="Kill running job by closing celery tasks and marking as killed",
        request=None,
        responses={
            204: None,
        },
    )
    @action(detail=True, methods=["patch"])
    def kill(self, request, pk=None):
        """
        Kill a running job by closing celery tasks and marking the job as killed.

        If the job is not running, raises a validation error.

        Returns:
        - No content (204) if the job is successfully killed.
        """
        # get job object or raise 404
        job = self.get_object()

        # check if job running
        if job.status in Job.Status.final_statuses():
            raise ValidationError({"detail": "Job is not running"})
        # close celery tasks and mark reports as killed
        job.kill_if_ongoing()
        return Response(status=status.HTTP_204_NO_CONTENT)

    @add_docs(
        description="Download file/sample associated with a job",
        request=None,
        responses={200: OpenApiTypes.BINARY, 400: None},
    )
    @action(detail=True, methods=["get"])
    def download_sample(self, request, pk=None):
        """
        Download a sample associated with a job.

        If the job does not have a sample, raises a validation error.

        Returns:
        - The file associated with the job as an attachment.

        :param url: pk (job_id)
        :returns: bytes
        """
        # get job object
        job = self.get_object()

        # make sure it is a sample
        if not job.is_sample:
            raise ValidationError(
                {"detail": "Requested job does not have a sample associated with it."}
            )
        return FileResponse(
            job.file,
            filename=job.file_name,
            content_type=job.file_mimetype,
            as_attachment=True,
        )

    @add_docs(description="Pivot a job")
    @action(
        detail=True, methods=["post"]
    )  # , url_path="pivot-(?P<pivot_config_pk>\d+)")
    def pivot(self, request, pk=None, pivot_config_pk=None):
        """
        Perform a pivot operation from a job's reports based on a specified pivot configuration.

        Expects the following parameters:
        - `pivot_config_pk`: The primary key of the pivot configuration to use.

        Returns:
        - List of job IDs created as a result of the pivot.
        """
        starting_job = self.get_object()
        try:
            pivot_config: PivotConfig = PivotConfig.objects.get(pk=pivot_config_pk)
        except PivotConfig.DoesNotExist:
            raise ValidationError({"detail": "Requested pivot config does not exist."})
        else:
            try:
                pivots = pivot_config.pivot_job(starting_job.reports)
            except KeyError:
                msg = (
                    f"Unable to retrieve value at {self.field}"
                    f" from job {starting_job.pk}"
                )
                logger.error(msg)
                raise ValidationError({"detail": msg})
            except Exception as e:
                logger.exception(e)
                raise ValidationError(
                    {"detail": f"Unable to start pivot from job {starting_job.pk}"}
                )
            else:
                return Response(
                    [pivot.ending_job.pk for pivot in pivots],
                    status=status.HTTP_201_CREATED,
                )

    @action(
        url_path="aggregate/status",
        detail=False,
        methods=["GET"],
    )
    @cache_action_response(timeout=60 * 5)
    def aggregate_status(self, request):
        """
        Aggregate jobs by their status.

        Returns:
        - Aggregated count of jobs for each status.
        """
        annotations = {
            key.lower(): Count("status", filter=Q(status=key))
            for key in Job.Status.values
        }
        return self.__aggregation_response_static(
            annotations, users=self.get_org_members(request)
        )

    @action(
        url_path="aggregate/type",
        detail=False,
        methods=["GET"],
    )
    @cache_action_response(timeout=60 * 5)
    def aggregate_type(self, request):
        """
        Aggregate jobs by type (file or observable).

        Returns:
        - Aggregated count of jobs for each type.
        """
        annotations = {
            "file": Count("is_sample", filter=Q(is_sample=True)),
            "observable": Count("is_sample", filter=Q(is_sample=False)),
        }
        return self.__aggregation_response_static(
            annotations, users=self.get_org_members(request)
        )

    @action(
        url_path="aggregate/observable_classification",
        detail=False,
        methods=["GET"],
    )
    @cache_action_response(timeout=60 * 5)
    def aggregate_observable_classification(self, request):
        """
        Aggregate jobs by observable classification.

        Returns:
        - Aggregated count of jobs for each observable classification.
        """
        annotations = {
            oc.lower(): Count(
                "observable_classification", filter=Q(observable_classification=oc)
            )
            for oc in ObservableTypes.values
        }
        return self.__aggregation_response_static(
            annotations, users=self.get_org_members(request)
        )

    @action(
        url_path="aggregate/file_mimetype",
        detail=False,
        methods=["GET"],
    )
    @cache_action_response(timeout=60 * 5)
    def aggregate_file_mimetype(self, request):
        """
        Aggregate jobs by file MIME type.

        Returns:
        - Aggregated count of jobs for each MIME type.
        """
        return self.__aggregation_response_dynamic(
            "file_mimetype", users=self.get_org_members(request)
        )

    @action(
        url_path="aggregate/observable_name",
        detail=False,
        methods=["GET"],
    )
    @cache_action_response(timeout=60 * 5)
    def aggregate_observable_name(self, request):
        """
        Aggregate jobs by observable name.

        Returns:
        - Aggregated count of jobs for each observable name.
        """
        return self.__aggregation_response_dynamic(
            "observable_name", False, users=self.get_org_members(request)
        )

    @action(
        url_path="aggregate/md5",
        detail=False,
        methods=["GET"],
    )
    @cache_action_response(timeout=60 * 5)
    def aggregate_md5(self, request):
        """
        Aggregate jobs by MD5 hash.

        Returns:
        - Aggregated count of jobs for each MD5 hash.
        """
        # this is for file
        return self.__aggregation_response_dynamic(
            "md5", False, users=self.get_org_members(request)
        )

    @staticmethod
    def get_org_members(request):
        """
        Retrieve members of the organization associated with the authenticated user.

        If the 'org' query parameter is set to 'true', this method returns all
        users who are members of the authenticated user's organization.

        Args:
            request: The HTTP request object containing user information and query parameters.

        Returns:
            list or None: A list of users who are members of the user's organization
            if the 'org' query parameter is 'true', otherwise None.
        """
        user = request.user
        org_param = request.GET.get("org", "").lower() == "true"
        users_of_organization = None
        if org_param:
            organization = user.membership.organization
            users_of_organization = [
                membership.user for membership in organization.members.all()
            ]
        return users_of_organization

    def __aggregation_response_static(self, annotations: dict, users=None) -> Response:
        """
        Generate a static aggregation of Job objects filtered by a time range.

        This method applies the provided annotations to aggregate Job objects
        within the specified time range. Optionally, it filters the results by
        the given list of users.

        Args:
            annotations (dict): Annotations to apply for the aggregation.
            users (list, optional): A list of users to filter the Job objects by.

        Returns:
            Response: A Django REST framework Response object containing the aggregated data.
        """
        delta, basis = self.__parse_range(self.request)
        filter_kwargs = {"received_request_time__gte": delta}
        if users:
            filter_kwargs["user__in"] = users
        qs = (
            Job.objects.filter(**filter_kwargs)
            .annotate(date=Trunc("received_request_time", basis))
            .values("date")
            .annotate(**annotations)
        )
        return Response(qs)

    def __aggregation_response_dynamic(
        self,
        field_name: str,
        group_by_date: bool = True,
        limit: int = 5,
        users=None,
    ) -> Response:
        """
        Dynamically aggregate Job objects based on a specified field and time range.

        This method identifies the most frequent values of a given field within
        a specified time range and aggregates the Job objects accordingly.
        Optionally, it can group the results by date and limit the number of
        most frequent values.

        Args:
            field_name (str): The name of the field to aggregate by.
            group_by_date (bool, optional): Whether to group the results by date. Defaults to True.
            limit (int, optional): The maximum number of most frequent values to retrieve. Defaults to 5.
            users (list, optional): A list of users to filter the Job objects by.

        Returns:
            Response: A Django REST framework Response object containing the most frequent values
            and the aggregated data.
        """
        delta, basis = self.__parse_range(self.request)
        filter_kwargs = {"received_request_time__gte": delta}
        if users:
            filter_kwargs["user__in"] = users
        if field_name == "md5":
            filter_kwargs["is_sample"] = True

        most_frequent_values = (
            Job.objects.filter(**filter_kwargs)
            .exclude(**{f"{field_name}__isnull": True})
            .exclude(**{f"{field_name}__exact": ""})
            # excluding those because they could lead to SQL query errors
            .exclude(
                observable_classification__in=[
                    ObservableClassification.URL,
                    ObservableClassification.GENERIC,
                ]
            )
            .annotate(count=Count(field_name))
            .distinct()
            .order_by("-count")[:limit]
            .values_list(field_name, flat=True)
        )

        logger.info(
            f"request: {field_name} found most_frequent_values: {most_frequent_values}"
        )

        if len(most_frequent_values):
            annotations = {
                val: Count(field_name, filter=Q(**{field_name: val}))
                for val in most_frequent_values
            }
            logger.debug(f"request: {field_name} annotations: {annotations}")
            if group_by_date:
                aggregation = (
                    Job.objects.filter(**filter_kwargs)
                    .annotate(date=Trunc("received_request_time", basis))
                    .values("date")
                    .annotate(**annotations)
                )
            else:
                aggregation = Job.objects.filter(**filter_kwargs).aggregate(
                    **annotations
                )
        else:
            aggregation = {}

        return Response(
            {
                "values": most_frequent_values,
                "aggregation": aggregation,
            }
        )

    @staticmethod
    def __parse_range(request):
        """
        Parse the time range from the request query parameters.

        This method attempts to extract the 'range' query parameter from the
        request. If the parameter is not provided, it defaults to '7d' (7 days).

        Args:
            request: The HTTP request object containing query parameters.

        Returns:
            tuple: A tuple containing the parsed time delta and the basis for date truncation.
        """
        try:
            range_str = request.GET["range"]
        except KeyError:
            # default
            range_str = "7d"

        return parse_humanized_range(range_str)