diff --git a/example.lua b/example.lua index fc6606e..96cc13d 100644 --- a/example.lua +++ b/example.lua @@ -20,7 +20,7 @@ end -- Download from http://dev.maxmind.com/geoip/geoip2/geolite2/ local mmdb = require "mmdb" -local geodb = assert(mmdb.read("GeoLite2-City.mmdb")) +local geodb = mmdb.read("GeoLite2-City.mmdb") if arg[1] then pp(geodb:search_ipv4(arg[1])) diff --git a/geo.lua b/geo.lua new file mode 100644 index 0000000..26429df --- /dev/null +++ b/geo.lua @@ -0,0 +1,50 @@ +#!/usr/bin/env ujit + +local argparse = require "argparse" +local mmdb = require "mmdb" +local json = require "dkjson" + +local function get_args() + local parser = argparse("geo", "Get info from MMDB database") + + parser:argument("ip", "IP v4/v6 address"):args "?" + + parser:flag("-i --info", "Infromation about database") + + parser + :option("-d --mmdb", "MaxMind2 database", "/usr/share/GeoIP/GeoIP2-City.mmdb") + + return parser:parse() +end + +local args = get_args() + +local geodb, err = mmdb.read(args.mmdb, true) +if not geodb then + io.stderr:write(err) + os.exit(-1) +end + +if args.info then + io.stdout:write(string.format( + "%s (%s); Support IPv%d; Nodes: %d\n", + geodb.data.database_type, + os.date("%F %T%z", tonumber(geodb.data.build_epoch)), + geodb.data.ip_version, + geodb.data.node_count + )) +end + +if args.ip then + local info + if string.find(args.ip, ':') then + info = geodb:search_ipv6(args.ip) + else + info = geodb:search_ipv4(args.ip) + end + if not info then + io.stderr:write('Not found info: ' .. args.ip) + os.exit(-2) + end + io.stdout:write(json.encode(info, { indent = true })) +end diff --git a/mmdb/init.lua b/mmdb/init.lua index 0d3b92f..614e4cb 100644 --- a/mmdb/init.lua +++ b/mmdb/init.lua @@ -14,7 +14,18 @@ local geodb_mt = { local data_types = {} local getters = {} -local function new(contents) +local function fail(safe, ...) + if safe then + return nil, ... + end + error(..., 2) +end + +local function new(contents, safe) + if not contents then + return fail(safe, "No MaxMind Database content") + end + local start_metadata do -- Find data section seperator; at most it's 128kb from the end local init = math.max(1, #contents-(128*1024)) @@ -24,11 +35,12 @@ local function new(contents) start_metadata = e + 1 end if start_metadata == nil then - return nil, "Invalid MaxMind Database" + return fail(safe, "Invalid MaxMind Database") end end local self = setmetatable({ + safe = not not safe; contents = contents; start_metadata = start_metadata; data = nil; @@ -42,7 +54,7 @@ local function new(contents) local getter = getters[data.record_size] if getter == nil then - return nil, "Unsupported record size: " .. data.record_size + return fail(self.safe, "Unsupported record size: " .. tostring(data.record_size)) end self.left, self.right, self.record_length = getter.left, getter.right, getter.record_length @@ -55,22 +67,22 @@ local function new(contents) return self end -local function read(filename) +local function read(filename, safe) local fd, err, errno = io.open(filename, "rb") if not fd then - return nil, err, errno + return fail(safe, err, errno) end local contents, err2, errno2 = fd:read("*a") fd:close() if not contents then - return nil, err2, errno2 + return fail(safe, err2, errno2) end - return new(contents) + return new(contents, safe) end -- Deprecated: local function open(filename) - return assert(read(filename)) + return read(filename, false) end function geodb_methods:read_data(base, offset) @@ -390,12 +402,17 @@ end local function ipv4_to_bit_array(str) local o1, o2, o3, o4 = str:match("(%d%d?%d?)%.(%d%d?%d?)%.(%d%d?%d?)%.(%d%d?%d?)") - assert(o1, "invalid IPv4 address") + if not o1 then + return nil, "invalid IPv4 address" + end o1 = tonumber(o1, 10) o2 = tonumber(o2, 10) o3 = tonumber(o3, 10) o4 = tonumber(o4, 10) - assert(o1 <= 255 and o2 <= 255 and o3 <= 255 and o4 <= 255, "invalid IPv4 address") + if not (o1 <= 255 and o2 <= 255 and o3 <= 255 and o4 <= 255) then + return nil, "invalid IPv4 address" + end + return { math.floor(o1 / 128) % 2 == 1; math.floor(o1 / 64) % 2 == 1; @@ -433,7 +450,11 @@ local function ipv4_to_bit_array(str) end function geodb_methods:search_ipv4(str) - return select(2, self:search(ipv4_to_bit_array(str), self.ipv4_start)) + local bits, err = ipv4_to_bit_array(str) + if not bits then + return fail(self.safe, err) + end + return select(2, self:search(bits, self.ipv4_start)) end local function ipv6_split(str) @@ -442,7 +463,9 @@ local function ipv6_split(str) for u16 in str:gmatch("(%x%x?%x?%x?):?") do n = n + 1 u16 = tonumber(u16, 16) - assert(u16, "invalid IPv6 address") + if not u16 then + return nil, "invalid IPv6 address" + end components[n] = u16 end return components, n @@ -451,9 +474,17 @@ end local function ipv6_to_bit_array(str) local a, b = str:match("^([%x:]-)::([%x:]*)$") local components, n = ipv6_split(a or str) + if not components then + return nil, n + end if a ~= nil then local end_components, m = ipv6_split(b) - assert(m+n <= 7, "invalid IPv6 address") + if not end_components then + return nil, m + end + if m+n > 7 then + return nil, "invalid IPv6 address" + end for i = n+1, 8-m do components[i] = 0 end @@ -461,7 +492,9 @@ local function ipv6_to_bit_array(str) components[i] = end_components[i-8+m] end else - assert(n == 8, "invalid IPv6 address") + if n ~= 8 then + return nil, "invalid IPv6 address" + end end -- Now components is an array of 16bit components local bits = {} @@ -475,7 +508,11 @@ local function ipv6_to_bit_array(str) end function geodb_methods:search_ipv6(str) - return select(2, self:search(ipv6_to_bit_array(str))) + local bits, err = ipv6_to_bit_array(str) + if not bits then + return fail(self.safe, err) + end + return select(2, self:search(bits)) end return {