#!/usr/bin/env python
"""Cron Job objects that get stored in the relational db."""

import logging

from grr.lib import rdfvalue
from grr.lib import stats
from grr.lib import utils
from grr.server.grr_response_server import aff4
from grr.server.grr_response_server import data_store
from grr.server.grr_response_server import flow
from grr.server.grr_response_server import queue_manager
from grr.server.grr_response_server.rdfvalues import cronjobs as rdf_cronjobs


class Error(Exception):
  pass


class LockError(Error):
  pass


class CronManager(object):
  """CronManager is used to schedule/terminate cron jobs."""

  def CreateJob(self, cron_args=None, job_id=None, disabled=False, token=None):
    """Creates a cron job that runs given flow with a given frequency.

    Args:
      cron_args: A protobuf of type CreateCronJobFlowArgs.

      job_id: Use this job_id instead of an autogenerated unique name (used
              for system cron jobs - we want them to have well-defined
              persistent name).

      disabled: If True, the job object will be created, but will be disabled.

      token: Security token used for data store access. Unused.

    Returns:
      URN of the cron job created.
    """
    # TODO(amoser): Remove the token from this method once the aff4
    # cronjobs are gone.
    del token
    if not job_id:
      uid = utils.PRNG.GetUInt16()
      job_id = "%s_%s" % (cron_args.flow_runner_args.flow_name, uid)

    job = rdf_cronjobs.CronJob(
        job_id=job_id, cron_args=cron_args, disabled=disabled)
    data_store.REL_DB.WriteCronJob(job)

    return job_id

  def ListJobs(self, token=None):
    """Returns a list of ids of all currently running cron jobs."""
    del token
    return [job.job_id for job in data_store.REL_DB.ReadCronJobs()]

  def ReadJob(self, job_id, token=None):
    del token
    return data_store.REL_DB.ReadCronJob(job_id)

  def ReadJobs(self, token=None):
    """Returns a list of all currently running cron jobs."""
    del token
    return data_store.REL_DB.ReadCronJobs()

  def ReadJobRuns(self, job_id, token=None):
    runs_base = rdfvalue.RDFURN("aff4:/cron").Add(job_id)
    fd = aff4.FACTORY.Open(runs_base, token=token)
    return list(fd.OpenChildren())

  def EnableJob(self, job_id, token=None):
    """Enable cron job with the given id."""
    del token
    return data_store.REL_DB.EnableCronJob(job_id)

  def DisableJob(self, job_id, token=None):
    """Disable cron job with the given id."""
    del token
    return data_store.REL_DB.DisableCronJob(job_id)

  def DeleteJob(self, job_id, token=None):
    """Deletes cron job with the given URN."""
    del token
    return data_store.REL_DB.DeleteCronJob(job_id)

  def RunOnce(self, token=None, force=False, names=None):
    """Tries to lock and run cron jobs.

    Args:
      token: security token.
      force: If True, force a run.
      names: List of cron jobs to run.  If unset, run them all.
    """
    leased_jobs = data_store.REL_DB.LeaseCronJobs(
        cronjob_ids=names, lease_time=rdfvalue.Duration("10m"))
    if not leased_jobs:
      return

    for job in leased_jobs:
      job.token = token
      try:
        logging.info("Running cron job: %s", job.job_id)
        self.RunJob(job, force=force, token=token)
      except Exception as e:  # pylint: disable=broad-except
        logging.exception("Error processing cron job %s: %s", job.job_id, e)
        stats.STATS.IncrementCounter("cron_internal_error")

    data_store.REL_DB.ReturnLeasedCronJobs(leased_jobs)

  def RunJob(self, job, force=False, token=None):
    """Does the actual work of the Cron, if the job is due to run.

    Args:
      job: The cronjob rdfvalue that should be run. Must be leased.
      force: If True, the job will run even if JobDueToRun() returns False.
      token: A datastore token.

    Raises:
      LockError: if the object is not locked.
    """
    if not job.leased_until or job.leased_until < rdfvalue.RDFDatetime.Now():
      raise LockError("CronJob must be leased for Run() to be called.")

    self.TerminateExpiredRun(job, token=token)

    # If currently running flow has finished, update our state.
    runs_base = rdfvalue.RDFURN("aff4:/cron").Add(job.job_id)
    if job.current_run_id:
      current_run_urn = runs_base.Add("F:%X" % job.current_run_id)
      current_flow = aff4.FACTORY.Open(current_run_urn, token=token)
      runner = current_flow.GetRunner()
      if not runner.IsRunning():
        if runner.context.state == "ERROR":
          status = rdf_cronjobs.CronJobRunStatus.Status.ERROR
          stats.STATS.IncrementCounter("cron_job_failure", fields=[job.job_id])
        else:
          status = rdf_cronjobs.CronJobRunStatus.Status.OK
          elapsed = rdfvalue.RDFDatetime.Now() - job.last_run_time
          stats.STATS.RecordEvent(
              "cron_job_latency", elapsed.seconds, fields=[job.job_id])

        data_store.REL_DB.UpdateCronJob(
            job.job_id, last_run_status=status, current_run_id=None)

    if not force and not self.JobDueToRun(job):
      return

    # Make sure the flow is created with cron job as a parent folder.
    job.cron_args.flow_runner_args.base_session_id = runs_base

    flow_urn = flow.GRRFlow.StartFlow(
        runner_args=job.cron_args.flow_runner_args,
        args=job.cron_args.flow_args,
        token=token,
        sync=False)

    job.current_run_id = int(flow_urn.Basename()[2:], 16)
    data_store.REL_DB.UpdateCronJob(
        job.job_id,
        last_run_time=rdfvalue.RDFDatetime.Now(),
        current_run_id=job.current_run_id)

  def TerminateExpiredRun(self, job, token=None):
    """Terminates the current run if it has exceeded CRON_ARGS.lifetime.

    Args:
      job: The cron job whose run should be terminated.
      token: Datastore token.

    Returns:
      bool: True if the run was terminate.
    """
    if not self.JobIsRunning(job, token=token):
      return False

    lifetime = job.cron_args.lifetime
    if not lifetime:
      return False

    elapsed = rdfvalue.RDFDatetime.Now() - job.last_run_time

    if elapsed > lifetime.seconds:
      self.StopCurrentRun(job, token=token)
      stats.STATS.IncrementCounter("cron_job_timeout", fields=[job.job_id])
      stats.STATS.RecordEvent("cron_job_latency", elapsed, fields=[job.job_id])
      return True

    return False

  def JobIsRunning(self, job, token=None):
    """Returns True if there's a currently running iteration of this job."""
    if not job.current_run_id:
      return False

    runs_base = rdfvalue.RDFURN("aff4:/cron").Add(job.job_id)
    run_urn = runs_base.Add("F:%X" % job.current_run_id)
    try:
      current_flow = aff4.FACTORY.Open(
          urn=run_urn, aff4_type=flow.GRRFlow, token=token, mode="r")
      return current_flow.GetRunner().IsRunning()

    except aff4.InstantiationError:
      # This isn't a flow, something went really wrong, clear it out.
      logging.error("Unable to open cron job run: %s", run_urn)
      data_store.REL_DB.UpdateCronJob(job.job_id, current_run_id=None)
      job.current_run_id = None
      return False

  def StopCurrentRun(self,
                     job,
                     reason="Cron lifetime exceeded.",
                     force=True,
                     token=None):
    """Stops the currently active run if there is one."""
    if not job.current_run_id:
      return

    runs_base = rdfvalue.RDFURN("aff4:/cron").Add(job.job_id)
    current_run_urn = runs_base.Add("F:%X" % job.current_run_id)
    flow.GRRFlow.TerminateFlow(
        current_run_urn, reason=reason, force=force, token=token)

    status = rdf_cronjobs.CronJobRunStatus.Status.TIMEOUT
    data_store.REL_DB.UpdateCronJob(
        job.job_id, last_run_status=status, current_run_id=None)

  def JobDueToRun(self, job):
    """Determines if the given job is due for another run.

    Args:
      job: The cron job rdfvalue object.
    Returns:
      True if it is time to run based on the specified frequency.
    """
    if job.disabled:
      return False

    now = rdfvalue.RDFDatetime.Now()

    # Its time to run.
    if (job.last_run_time is None or
        now > job.cron_args.periodicity.Expiry(job.last_run_time)):

      # Not due to start yet.
      if now < job.cron_args.start_time:
        return False

      # Do we allow overruns?
      if job.cron_args.allow_overruns:
        return True

      # No currently executing job - lets go.
      if not job.current_run_id:
        return True

    return False

  def DeleteRuns(self, job, age=None, token=None):
    """Deletes runs that were started before the age given."""
    if age is None:
      raise ValueError("age can't be None")

    runs_base = rdfvalue.RDFURN("aff4:/cron").Add(job.job_id)
    runs_base_obj = aff4.FACTORY.Open(runs_base, token=token)
    child_flows = list(runs_base_obj.ListChildren(age=age))
    with queue_manager.QueueManager(token=token) as queuemanager:
      queuemanager.MultiDestroyFlowStates(child_flows)

    aff4.FACTORY.MultiDelete(child_flows, token=token)
