1# (c) 2013, Jan-Piet Mens <jpmens(at)gmail.com>
2# (c) 2017 Ansible Project
3# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
4from __future__ import (absolute_import, division, print_function)
5__metaclass__ = type
6
7DOCUMENTATION = """
8    lookup: csvfile
9    author: Jan-Piet Mens (@jpmens) <jpmens(at)gmail.com>
10    version_added: "1.5"
11    short_description: read data from a TSV or CSV file
12    description:
13      - The csvfile lookup reads the contents of a file in CSV (comma-separated value) format.
14        The lookup looks for the row where the first column matches keyname, and returns the value in the second column, unless a different column is specified.
15    options:
16      col:
17        description:  column to return (0 index).
18        default: "1"
19      default:
20        description: what to return if the value is not found in the file.
21        default: ''
22      delimiter:
23        description: field separator in the file, for a tab you can specify "TAB" or "t".
24        default: TAB
25      file:
26        description: name of the CSV/TSV file to open.
27        default: ansible.csv
28      encoding:
29        description: Encoding (character set) of the used CSV file.
30        default: utf-8
31        version_added: "2.1"
32    notes:
33      - The default is for TSV files (tab delimited) not CSV (comma delimited) ... yes the name is misleading.
34"""
35
36EXAMPLES = """
37- name:  Match 'Li' on the first column, return the second column (0 based index)
38  debug: msg="The atomic number of Lithium is {{ lookup('csvfile', 'Li file=elements.csv delimiter=,') }}"
39
40- name: msg="Match 'Li' on the first column, but return the 3rd column (columns start counting after the match)"
41  debug: msg="The atomic mass of Lithium is {{ lookup('csvfile', 'Li file=elements.csv delimiter=, col=2') }}"
42
43- name: Define Values From CSV File
44  set_fact:
45    loop_ip: "{{ lookup('csvfile', bgp_neighbor_ip +' file=bgp_neighbors.csv delimiter=, col=1') }}"
46    int_ip: "{{ lookup('csvfile', bgp_neighbor_ip +' file=bgp_neighbors.csv delimiter=, col=2') }}"
47    int_mask: "{{ lookup('csvfile', bgp_neighbor_ip +' file=bgp_neighbors.csv delimiter=, col=3') }}"
48    int_name: "{{ lookup('csvfile', bgp_neighbor_ip +' file=bgp_neighbors.csv delimiter=, col=4') }}"
49    local_as: "{{ lookup('csvfile', bgp_neighbor_ip +' file=bgp_neighbors.csv delimiter=, col=5') }}"
50    neighbor_as: "{{ lookup('csvfile', bgp_neighbor_ip +' file=bgp_neighbors.csv delimiter=, col=6') }}"
51    neigh_int_ip: "{{ lookup('csvfile', bgp_neighbor_ip +' file=bgp_neighbors.csv delimiter=, col=7') }}"
52  delegate_to: localhost
53"""
54
55RETURN = """
56  _raw:
57    description:
58      - value(s) stored in file column
59"""
60
61import codecs
62import csv
63
64from ansible.errors import AnsibleError, AnsibleAssertionError
65from ansible.plugins.lookup import LookupBase
66from ansible.module_utils.six import PY2
67from ansible.module_utils._text import to_bytes, to_native, to_text
68from ansible.module_utils.common._collections_compat import MutableSequence
69
70
71class CSVRecoder:
72    """
73    Iterator that reads an encoded stream and reencodes the input to UTF-8
74    """
75    def __init__(self, f, encoding='utf-8'):
76        self.reader = codecs.getreader(encoding)(f)
77
78    def __iter__(self):
79        return self
80
81    def __next__(self):
82        return next(self.reader).encode("utf-8")
83
84    next = __next__   # For Python 2
85
86
87class CSVReader:
88    """
89    A CSV reader which will iterate over lines in the CSV file "f",
90    which is encoded in the given encoding.
91    """
92
93    def __init__(self, f, dialect=csv.excel, encoding='utf-8', **kwds):
94        if PY2:
95            f = CSVRecoder(f, encoding)
96        else:
97            f = codecs.getreader(encoding)(f)
98
99        self.reader = csv.reader(f, dialect=dialect, **kwds)
100
101    def __next__(self):
102        row = next(self.reader)
103        return [to_text(s) for s in row]
104
105    next = __next__  # For Python 2
106
107    def __iter__(self):
108        return self
109
110
111class LookupModule(LookupBase):
112
113    def read_csv(self, filename, key, delimiter, encoding='utf-8', dflt=None, col=1):
114
115        try:
116            f = open(filename, 'rb')
117            creader = CSVReader(f, delimiter=to_native(delimiter), encoding=encoding)
118
119            for row in creader:
120                if len(row) and row[0] == key:
121                    return row[int(col)]
122        except Exception as e:
123            raise AnsibleError("csvfile: %s" % to_native(e))
124
125        return dflt
126
127    def run(self, terms, variables=None, **kwargs):
128
129        ret = []
130
131        for term in terms:
132            params = term.split()
133            key = params[0]
134
135            paramvals = {
136                'col': "1",          # column to return
137                'default': None,
138                'delimiter': "TAB",
139                'file': 'ansible.csv',
140                'encoding': 'utf-8',
141            }
142
143            # parameters specified?
144            try:
145                for param in params[1:]:
146                    name, value = param.split('=')
147                    if name not in paramvals:
148                        raise AnsibleAssertionError('%s not in paramvals' % name)
149                    paramvals[name] = value
150            except (ValueError, AssertionError) as e:
151                raise AnsibleError(e)
152
153            if paramvals['delimiter'] == 'TAB':
154                paramvals['delimiter'] = "\t"
155
156            lookupfile = self.find_file_in_search_path(variables, 'files', paramvals['file'])
157            var = self.read_csv(lookupfile, key, paramvals['delimiter'], paramvals['encoding'], paramvals['default'], paramvals['col'])
158            if var is not None:
159                if isinstance(var, MutableSequence):
160                    for v in var:
161                        ret.append(v)
162                else:
163                    ret.append(var)
164        return ret
165