1# (c) 2020 Ansible Project
2# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
3from __future__ import absolute_import, division, print_function
4
5__metaclass__ = type
6
7DOCUMENTATION = """
8    lookup: schedule_rrule
9    author: John Westcott IV (@john-westcott-iv)
10    short_description: Generate an rrule string which can be used for Schedules
11    requirements:
12      - pytz
13      - python-dateutil >= 2.7.0
14    description:
15      - Returns a string based on criteria which represents an rrule
16    options:
17      _terms:
18        description:
19          - The frequency of the schedule
20          - none - Run this schedule once
21          - minute - Run this schedule every x minutes
22          - hour - Run this schedule every x hours
23          - day - Run this schedule every x days
24          - week - Run this schedule weekly
25          - month - Run this schedule monthly
26        required: True
27        choices: ['none', 'minute', 'hour', 'day', 'week', 'month']
28      start_date:
29        description:
30          - The date to start the rule
31          - Used for all frequencies
32          - Format should be YYYY-MM-DD [HH:MM:SS]
33        type: str
34      timezone:
35        description:
36          - The timezone to use for this rule
37          - Used for all frequencies
38          - Format should be as US/Eastern
39          - Defaults to America/New_York
40        type: str
41      every:
42        description:
43          - The repetition in months, weeks, days hours or minutes
44          - Used for all types except none
45        type: int
46      end_on:
47        description:
48          - How to end this schedule
49          - If this is not defined, this schedule will never end
50          - If this is a positive integer, this schedule will end after this number of occurences
51          - If this is a date in the format YYYY-MM-DD [HH:MM:SS], this schedule ends after this date
52          - Used for all types except none
53        type: str
54      on_days:
55        description:
56          - The days to run this schedule on
57          - A comma-separated list which can contain values sunday, monday, tuesday, wednesday, thursday, friday
58          - Used for week type schedules
59      month_day_number:
60        description:
61          - The day of the month this schedule will run on (0-31)
62          - Used for month type schedules
63          - Cannot be used with on_the parameter
64        type: int
65      on_the:
66        description:
67          - A description on when this schedule will run
68          - Two strings separated by a space
69          - First string is one of first, second, third, fourth, last
70          - Second string is one of sunday, monday, tuesday, wednesday, thursday, friday
71          - Used for month type schedules
72          - Cannot be used with month_day_number parameters
73"""
74
75EXAMPLES = """
76    - name: Create a string for a schedule
77      debug:
78        msg: "{{ query('awx.awx.schedule_rrule', 'none', start_date='1979-09-13 03:45:07') }}"
79"""
80
81RETURN = """
82_raw:
83  description:
84    - String in the rrule format
85  type: string
86"""
87import re
88
89from ansible.module_utils.six import raise_from
90from ansible.plugins.lookup import LookupBase
91from ansible.errors import AnsibleError
92from datetime import datetime
93from dateutil import rrule
94
95try:
96    import pytz
97    from dateutil import rrule
98except ImportError as imp_exc:
99    LIBRARY_IMPORT_ERROR = imp_exc
100else:
101    LIBRARY_IMPORT_ERROR = None
102
103
104class LookupModule(LookupBase):
105    frequencies = {
106        'none': rrule.DAILY,
107        'minute': rrule.MINUTELY,
108        'hour': rrule.HOURLY,
109        'day': rrule.DAILY,
110        'week': rrule.WEEKLY,
111        'month': rrule.MONTHLY,
112    }
113
114    weekdays = {
115        'monday': rrule.MO,
116        'tuesday': rrule.TU,
117        'wednesday': rrule.WE,
118        'thursday': rrule.TH,
119        'friday': rrule.FR,
120        'saturday': rrule.SA,
121        'sunday': rrule.SU,
122    }
123
124    set_positions = {
125        'first': 1,
126        'second': 2,
127        'third': 3,
128        'fourth': 4,
129        'last': -1,
130    }
131
132    # plugin constructor
133    def __init__(self, *args, **kwargs):
134        if LIBRARY_IMPORT_ERROR:
135            raise_from(
136                AnsibleError('{0}'.format(LIBRARY_IMPORT_ERROR)),
137                LIBRARY_IMPORT_ERROR
138            )
139        super().__init__(*args, **kwargs)
140
141    @staticmethod
142    def parse_date_time(date_string):
143        try:
144            return datetime.strptime(date_string, '%Y-%m-%d %H:%M:%S')
145        except ValueError:
146            return datetime.strptime(date_string, '%Y-%m-%d')
147
148    def run(self, terms, variables=None, **kwargs):
149        if len(terms) != 1:
150            raise AnsibleError('You may only pass one schedule type in at a time')
151
152        frequency = terms[0].lower()
153
154        return self.get_rrule(frequency, kwargs)
155
156    @staticmethod
157    def get_rrule(frequency, kwargs):
158
159        if frequency not in LookupModule.frequencies:
160            raise AnsibleError('Frequency of {0} is invalid'.format(frequency))
161
162        rrule_kwargs = {
163            'freq': LookupModule.frequencies[frequency],
164            'interval': kwargs.get('every', 1),
165        }
166
167        # All frequencies can use a start date
168        if 'start_date' in kwargs:
169            try:
170                rrule_kwargs['dtstart'] = LookupModule.parse_date_time(kwargs['start_date'])
171            except Exception as e:
172                raise_from(AnsibleError('Parameter start_date must be in the format YYYY-MM-DD [HH:MM:SS]'), e)
173
174        # If we are a none frequency we don't need anything else
175        if frequency == 'none':
176            rrule_kwargs['count'] = 1
177        else:
178            # All non-none frequencies can have an end_on option
179            if 'end_on' in kwargs:
180                end_on = kwargs['end_on']
181                if re.match(r'^\d+$', end_on):
182                    rrule_kwargs['count'] = end_on
183                else:
184                    try:
185                        rrule_kwargs['until'] = LookupModule.parse_date_time(end_on)
186                    except Exception as e:
187                        raise_from(AnsibleError('Parameter end_on must either be an integer or in the format YYYY-MM-DD [HH:MM:SS]'), e)
188
189            # A week-based frequency can also take the on_days parameter
190            if frequency == 'week' and 'on_days' in kwargs:
191                days = []
192                for day in kwargs['on_days'].split(','):
193                    day = day.strip()
194                    if day not in LookupModule.weekdays:
195                        raise AnsibleError('Parameter on_days must only contain values {0}'.format(', '.join(LookupModule.weekdays.keys())))
196                    days.append(LookupModule.weekdays[day])
197
198                rrule_kwargs['byweekday'] = days
199
200            # A month-based frequency can also deal with month_day_number and on_the options
201            if frequency == 'month':
202                if 'month_day_number' in kwargs and 'on_the' in kwargs:
203                    raise AnsibleError('Month based frequencies can have month_day_number or on_the but not both')
204
205                if 'month_day_number' in kwargs:
206                    try:
207                        my_month_day = int(kwargs['month_day_number'])
208                        if my_month_day < 1 or my_month_day > 31:
209                            raise Exception()
210                    except Exception as e:
211                        raise_from(AnsibleError('month_day_number must be between 1 and 31'), e)
212
213                    rrule_kwargs['bymonthday'] = my_month_day
214
215                if 'on_the' in kwargs:
216                    try:
217                        (occurance, weekday) = kwargs['on_the'].split(' ')
218                    except Exception as e:
219                        raise_from(AnsibleError('on_the parameter must be two words separated by a space'), e)
220
221                    if weekday not in LookupModule.weekdays:
222                        raise AnsibleError('Weekday portion of on_the parameter is not valid')
223                    if occurance not in LookupModule.set_positions:
224                        raise AnsibleError('The first string of the on_the parameter is not valid')
225
226                    rrule_kwargs['byweekday'] = LookupModule.weekdays[weekday]
227                    rrule_kwargs['bysetpos'] = LookupModule.set_positions[occurance]
228
229        my_rule = rrule.rrule(**rrule_kwargs)
230
231        # All frequencies can use a timezone but rrule can't support the format that AWX uses.
232        # So we will do a string manip here if we need to
233        timezone = 'America/New_York'
234        if 'timezone' in kwargs:
235            if kwargs['timezone'] not in pytz.all_timezones:
236                raise AnsibleError('Timezone parameter is not valid')
237            timezone = kwargs['timezone']
238
239        # rrule puts a \n in the rule instad of a space and can't handle timezones
240        return_rrule = str(my_rule).replace('\n', ' ').replace('DTSTART:', 'DTSTART;TZID={0}:'.format(timezone))
241        # AWX requires an interval. rrule will not add interval if it's set to 1
242        if kwargs.get('every', 1) == 1:
243            return_rrule = "{0};INTERVAL=1".format(return_rrule)
244
245        return return_rrule
246