|
| 1 | +-- This Source Code Form is subject to the terms of the Mozilla Public |
| 2 | +-- License, v. 2.0. If a copy of the MPL was not distributed with this |
| 3 | +-- file, You can obtain one at http://mozilla.org/MPL/2.0/. |
| 4 | + |
| 5 | +--[[ |
| 6 | +# Heka Message Extensions with GeoIP information |
| 7 | +
|
| 8 | +This module is intended to be used with another IO module such as a decoder to |
| 9 | +extend the message with GeoIP information prior to injection. |
| 10 | +
|
| 11 | +## Functions |
| 12 | +
|
| 13 | +### add_geoip |
| 14 | +
|
| 15 | +Given a Heka message add the geoip entries based on the specified field name. |
| 16 | +The requested geoip entries are specified in the `lookup` configuration table. |
| 17 | +The table consists of the entry name and the suffix to add to the field name, |
| 18 | +an empty suffix will overwrite the original value. See: the `opts` variable in |
| 19 | +https://github.com/agladysh/lua-geoip/blob/master/src/city.c for the available |
| 20 | +options, common values include "city" and "country_code". |
| 21 | +
|
| 22 | +For example, if "remote_addr" is present in msg.Fields with a cfg of `lookup = |
| 23 | +{ city = "_city" }` on a successful lookup a new field will be added to the |
| 24 | +message named "remote_addr_city" with the resulting city value. |
| 25 | +
|
| 26 | +*Arguments* |
| 27 | +- msg (table) - original message |
| 28 | +- field_name (string) - field name in the message to lookup |
| 29 | +
|
| 30 | +*Return* |
| 31 | +- none - the message is modified in place or an error is thrown |
| 32 | +
|
| 33 | +## Configuration examples |
| 34 | +```lua |
| 35 | +geoip_heka = { |
| 36 | + city_db_file = "/path/to/geo/dat", -- path to GeoIP data |
| 37 | + lookup = { city = "_city", country_code = "_country" }, -- entries to lookup and their field suffix |
| 38 | +
|
| 39 | + test = false, -- true if being used in tests without GeoIP database |
| 40 | +} |
| 41 | +``` |
| 42 | +--]] |
| 43 | +local module_name = ... |
| 44 | +local module_cfg = require "string".gsub(module_name, "%.", "_") |
| 45 | +local cfg = read_config(module_cfg) or error(module_name .. " configuration not found") |
| 46 | +assert(type(cfg.lookup) == "table" and next(cfg.lookup) ~= nil, "lookup configuration must be a table") |
| 47 | + |
| 48 | +local geoip = require "geoip.city" |
| 49 | +local ostime = require "os".time |
| 50 | +local smatch = string.match |
| 51 | +local assert = assert |
| 52 | +local pairs = pairs |
| 53 | +local type = type |
| 54 | + |
| 55 | +local M = {} |
| 56 | +setfenv(1, M) -- Remove external access to contain everything in the module. |
| 57 | + |
| 58 | +local gtest = {} |
| 59 | + |
| 60 | +-- In test mode, returns Mountain View, US for 192.168.1.2 otherwise nil, only |
| 61 | +-- supports country_code and city lookup values |
| 62 | +local test_return = {country_code = "US", city = "Mountain View"} |
| 63 | +function gtest:query_by_addr(v, lookup) |
| 64 | + if v == "192.168.1.2" then |
| 65 | + return test_return[lookup] |
| 66 | + end |
| 67 | +end |
| 68 | + |
| 69 | +local ptime = 0 |
| 70 | +local geodb = nil |
| 71 | +local function refresh_db() |
| 72 | + local ctime = ostime() |
| 73 | + if ptime + 3600 < ctime then |
| 74 | + if geodb then geodb:close() end |
| 75 | + geodb = assert(geoip.open(cfg.city_db_file)) |
| 76 | + ptime = ctime |
| 77 | + end |
| 78 | +end |
| 79 | + |
| 80 | + |
| 81 | +local function validate_ip(ip) |
| 82 | + local sl = #ip |
| 83 | + return not (sl < 7 or sl > 15 or not smatch(ip, "^%d+%.%d+%.%d+%.%d+$")) |
| 84 | +end |
| 85 | + |
| 86 | + |
| 87 | +if cfg.test then |
| 88 | + geodb = gtest |
| 89 | + refresh_db = function () return end |
| 90 | +else |
| 91 | + refresh_db() |
| 92 | +end |
| 93 | + |
| 94 | +function add_geoip(msg, field_name) |
| 95 | + if not msg.Fields then return end |
| 96 | + local value = msg.Fields[field_name] |
| 97 | + if type(value) ~= "string" or not validate_ip(value) then return end |
| 98 | + |
| 99 | + refresh_db() |
| 100 | + |
| 101 | + for k,v in pairs(cfg.lookup) do |
| 102 | + local ret = geodb:query_by_addr(value, k) |
| 103 | + if ret then |
| 104 | + msg.Fields[field_name .. v] = ret |
| 105 | + if cfg.remove_original_field then |
| 106 | + msg.Fields[field_name] = nil |
| 107 | + end |
| 108 | + end |
| 109 | + end |
| 110 | +end |
| 111 | + |
| 112 | + |
| 113 | +return M |
0 commit comments