1"""
2Wrapper module for at(1)
3
4Also, a 'tag' feature has been added to more
5easily tag jobs.
6
7:platform:      linux,openbsd,freebsd
8
9.. versionchanged:: 2017.7.0
10"""
11
12import datetime
13import re
14import time
15
16import salt.utils.data
17import salt.utils.path
18import salt.utils.platform
19
20# pylint: enable=import-error,redefined-builtin
21from salt.exceptions import CommandNotFoundError
22
23# pylint: disable=import-error,redefined-builtin
24
25# OS Families that should work (Ubuntu and Debian are the default)
26# TODO: Refactor some of this module to remove the checks for binaries
27
28# Tested on OpenBSD 5.0
29BSD = ("OpenBSD", "FreeBSD")
30
31__virtualname__ = "at"
32
33
34def __virtual__():
35    """
36    Most everything has the ability to support at(1)
37    """
38    if salt.utils.platform.is_windows() or salt.utils.platform.is_sunos():
39        return (False, "The at module could not be loaded: unsupported platform")
40    if salt.utils.path.which("at") is None:
41        return (False, "The at module could not be loaded: at command not found")
42    return __virtualname__
43
44
45def _cmd(binary, *args):
46    """
47    Wrapper to run at(1) or return None.
48    """
49    binary = salt.utils.path.which(binary)
50    if not binary:
51        raise CommandNotFoundError("{}: command not found".format(binary))
52    cmd = [binary] + list(args)
53    return __salt__["cmd.run_stdout"]([binary] + list(args), python_shell=False)
54
55
56def atq(tag=None):
57    """
58    List all queued and running jobs or only those with
59    an optional 'tag'.
60
61    CLI Example:
62
63    .. code-block:: bash
64
65        salt '*' at.atq
66        salt '*' at.atq [tag]
67        salt '*' at.atq [job number]
68    """
69    jobs = []
70
71    # Shim to produce output similar to what __virtual__() should do
72    # but __salt__ isn't available in __virtual__()
73    # Tested on CentOS 5.8
74    if __grains__["os_family"] == "RedHat":
75        output = _cmd("at", "-l")
76    else:
77        output = _cmd("atq")
78
79    if output is None:
80        return "'at.atq' is not available."
81
82    # No jobs so return
83    if output == "":
84        return {"jobs": jobs}
85
86    # Jobs created with at.at() will use the following
87    # comment to denote a tagged job.
88    job_kw_regex = re.compile(r"^### SALT: (\w+)")
89
90    # Split each job into a dictionary and handle
91    # pulling out tags or only listing jobs with a certain
92    # tag
93    for line in output.splitlines():
94        job_tag = ""
95
96        # Redhat/CentOS
97        if __grains__["os_family"] == "RedHat":
98            job, spec = line.split("\t")
99            specs = spec.split()
100        elif __grains__["os"] == "OpenBSD":
101            if line.startswith(" Rank"):
102                continue
103            else:
104                tmp = line.split()
105                timestr = " ".join(tmp[1:5])
106                job = tmp[6]
107                specs = (
108                    datetime.datetime(*(time.strptime(timestr, "%b %d, %Y %H:%M")[0:5]))
109                    .isoformat()
110                    .split("T")
111                )
112                specs.append(tmp[7])
113                specs.append(tmp[5])
114        elif __grains__["os"] == "FreeBSD":
115            if line.startswith("Date"):
116                continue
117            else:
118                tmp = line.split()
119                timestr = " ".join(tmp[1:6])
120                job = tmp[8]
121                specs = (
122                    datetime.datetime(
123                        *(time.strptime(timestr, "%b %d %H:%M:%S %Z %Y")[0:5])
124                    )
125                    .isoformat()
126                    .split("T")
127                )
128                specs.append(tmp[7])
129                specs.append(tmp[6])
130
131        else:
132            job, spec = line.split("\t")
133            tmp = spec.split()
134            timestr = " ".join(tmp[0:5])
135            specs = (
136                datetime.datetime(*(time.strptime(timestr)[0:5])).isoformat().split("T")
137            )
138            specs.append(tmp[5])
139            specs.append(tmp[6])
140
141        # Search for any tags
142        atc_out = _cmd("at", "-c", job)
143        for line in atc_out.splitlines():
144            tmp = job_kw_regex.match(line)
145            if tmp:
146                job_tag = tmp.groups()[0]
147
148        if __grains__["os"] in BSD:
149            job = str(job)
150        else:
151            job = int(job)
152
153        # If a tag is supplied, only list jobs with that tag
154        if tag:
155            # TODO: Looks like there is a difference between salt and salt-call
156            # If I don't wrap job in an int(), it fails on salt but works on
157            # salt-call. With the int(), it fails with salt-call but not salt.
158            if tag == job_tag or tag == job:
159                jobs.append(
160                    {
161                        "job": job,
162                        "date": specs[0],
163                        "time": specs[1],
164                        "queue": specs[2],
165                        "user": specs[3],
166                        "tag": job_tag,
167                    }
168                )
169        else:
170            jobs.append(
171                {
172                    "job": job,
173                    "date": specs[0],
174                    "time": specs[1],
175                    "queue": specs[2],
176                    "user": specs[3],
177                    "tag": job_tag,
178                }
179            )
180
181    return {"jobs": jobs}
182
183
184def atrm(*args):
185    """
186    Remove jobs from the queue.
187
188    CLI Example:
189
190    .. code-block:: bash
191
192        salt '*' at.atrm <jobid> <jobid> .. <jobid>
193        salt '*' at.atrm all
194        salt '*' at.atrm all [tag]
195    """
196
197    # Need to do this here also since we use atq()
198    if not salt.utils.path.which("at"):
199        return "'at.atrm' is not available."
200
201    if not args:
202        return {"jobs": {"removed": [], "tag": None}}
203
204    # Convert all to strings
205    args = salt.utils.data.stringify(args)
206
207    if args[0] == "all":
208        if len(args) > 1:
209            opts = list(list(map(str, [j["job"] for j in atq(args[1])["jobs"]])))
210            ret = {"jobs": {"removed": opts, "tag": args[1]}}
211        else:
212            opts = list(list(map(str, [j["job"] for j in atq()["jobs"]])))
213            ret = {"jobs": {"removed": opts, "tag": None}}
214    else:
215        opts = list(
216            list(
217                map(
218                    str,
219                    [i["job"] for i in atq()["jobs"] if str(i["job"]) in args],
220                )
221            )
222        )
223        ret = {"jobs": {"removed": opts, "tag": None}}
224
225    # Shim to produce output similar to what __virtual__() should do
226    # but __salt__ isn't available in __virtual__()
227    output = _cmd("at", "-d", " ".join(opts))
228    if output is None:
229        return "'at.atrm' is not available."
230
231    return ret
232
233
234def at(*args, **kwargs):  # pylint: disable=C0103
235    """
236    Add a job to the queue.
237
238    The 'timespec' follows the format documented in the
239    at(1) manpage.
240
241    CLI Example:
242
243    .. code-block:: bash
244
245        salt '*' at.at <timespec> <cmd> [tag=<tag>] [runas=<user>]
246        salt '*' at.at 12:05am '/sbin/reboot' tag=reboot
247        salt '*' at.at '3:05am +3 days' 'bin/myscript' tag=nightly runas=jim
248    """
249
250    if len(args) < 2:
251        return {"jobs": []}
252
253    # Shim to produce output similar to what __virtual__() should do
254    # but __salt__ isn't available in __virtual__()
255    binary = salt.utils.path.which("at")
256    if not binary:
257        return "'at.at' is not available."
258
259    if "tag" in kwargs:
260        stdin = "### SALT: {}\n{}".format(kwargs["tag"], " ".join(args[1:]))
261    else:
262        stdin = " ".join(args[1:])
263    cmd = [binary, args[0]]
264
265    cmd_kwargs = {"stdin": stdin, "python_shell": False}
266    if "runas" in kwargs:
267        cmd_kwargs["runas"] = kwargs["runas"]
268    output = __salt__["cmd.run"](cmd, **cmd_kwargs)
269
270    if output is None:
271        return "'at.at' is not available."
272
273    if output.endswith("Garbled time"):
274        return {"jobs": [], "error": "invalid timespec"}
275
276    if output.startswith("warning: commands"):
277        output = output.splitlines()[1]
278
279    if output.startswith("commands will be executed"):
280        output = output.splitlines()[1]
281
282    output = output.split()[1]
283
284    if __grains__["os"] in BSD:
285        return atq(str(output))
286    else:
287        return atq(int(output))
288
289
290def atc(jobid):
291    """
292    Print the at(1) script that will run for the passed job
293    id. This is mostly for debugging so the output will
294    just be text.
295
296    CLI Example:
297
298    .. code-block:: bash
299
300        salt '*' at.atc <jobid>
301    """
302    # Shim to produce output similar to what __virtual__() should do
303    # but __salt__ isn't available in __virtual__()
304    output = _cmd("at", "-c", str(jobid))
305
306    if output is None:
307        return "'at.atc' is not available."
308    elif output == "":
309        return {"error": "invalid job id '{}'".format(jobid)}
310
311    return output
312
313
314def _atq(**kwargs):
315    """
316    Return match jobs list
317    """
318
319    jobs = []
320
321    runas = kwargs.get("runas", None)
322    tag = kwargs.get("tag", None)
323    hour = kwargs.get("hour", None)
324    minute = kwargs.get("minute", None)
325    day = kwargs.get("day", None)
326    month = kwargs.get("month", None)
327    year = kwargs.get("year", None)
328    if year and len(str(year)) == 2:
329        year = "20{}".format(year)
330
331    jobinfo = atq()["jobs"]
332    if not jobinfo:
333        return {"jobs": jobs}
334
335    for job in jobinfo:
336
337        if not runas:
338            pass
339        elif runas == job["user"]:
340            pass
341        else:
342            continue
343
344        if not tag:
345            pass
346        elif tag == job["tag"]:
347            pass
348        else:
349            continue
350
351        if not hour:
352            pass
353        elif "{:02d}".format(int(hour)) == job["time"].split(":")[0]:
354            pass
355        else:
356            continue
357
358        if not minute:
359            pass
360        elif "{:02d}".format(int(minute)) == job["time"].split(":")[1]:
361            pass
362        else:
363            continue
364
365        if not day:
366            pass
367        elif "{:02d}".format(int(day)) == job["date"].split("-")[2]:
368            pass
369        else:
370            continue
371
372        if not month:
373            pass
374        elif "{:02d}".format(int(month)) == job["date"].split("-")[1]:
375            pass
376        else:
377            continue
378
379        if not year:
380            pass
381        elif year == job["date"].split("-")[0]:
382            pass
383        else:
384            continue
385
386        jobs.append(job)
387
388    if not jobs:
389        note = "No match jobs or time format error"
390        return {"jobs": jobs, "note": note}
391
392    return {"jobs": jobs}
393
394
395def jobcheck(**kwargs):
396    """
397    Check the job from queue.
398    The kwargs dict include 'hour minute day month year tag runas'
399    Other parameters will be ignored.
400
401    CLI Example:
402
403    .. code-block:: bash
404
405        salt '*' at.jobcheck runas=jam day=13
406        salt '*' at.jobcheck day=13 month=12 year=13 tag=rose
407    """
408
409    if not kwargs:
410        return {"error": "You have given a condition"}
411
412    return _atq(**kwargs)
413