{-# LANGUAGE CPP #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE RankNTypes #-}
module Data.GeoIP2 (
GeoDB
, openGeoDB, openGeoDBBS
, geoDbLanguages, geoDbType, geoDbDescription
, geoDbAddrType, GeoIP(..)
, DecodeException(..)
, findGeoData
, GeoResult(..)
, Location(..)
, AS(..)
, GeoField, GeoFieldT(..)
, rawGeoData
, _DataString, _DataDouble, _DataInt, _DataWord
, _DataMap, _DataArray, _DataBool, _DataUnknown
, key
, geoNum
) where
#if !MIN_VERSION_base(4,8,0)
import Control.Applicative ((<$>), (<*>))
#endif
import Control.Monad (unless, when)
import qualified Data.ByteString as BS
import Data.Int
import Data.IP (IP (..), ipv4ToIPv6)
import Data.Maybe (mapMaybe)
import Data.Serialize
import qualified Data.Text as T
import System.IO.MMap
import Control.Lens (ix, (^?), _Just, to, (^..), Traversal', Fold, prism')
import Control.Exception (throwIO, Exception)
import Data.GeoIP2.Fields
import Data.GeoIP2.SearchTree
data GeoIP = GeoIPv6 | GeoIPv4 deriving (Eq, Show)
data DecodeException = DecodeException String
deriving (Show)
instance Exception DecodeException
data GeoDB = GeoDB {
geoMem :: BS.ByteString
, geoDbType :: T.Text
, geoDbLanguages :: [T.Text]
, geoDbNodeCount :: Int64
, geoDbRecordSize :: Int
, geoDbAddrType :: GeoIP
, geoDbDescription :: Maybe T.Text
}
getHeaderBytes :: BS.ByteString -> BS.ByteString
getHeaderBytes = lastsubstring "\xab\xcd\xefMaxMind.com"
where
lastsubstring pattern string =
case BS.breakSubstring pattern string of
(res, "") -> res
(_, rest) -> lastsubstring pattern (BS.drop (BS.length pattern) rest)
openGeoDB :: FilePath -> IO GeoDB
openGeoDB geoFile = do
bsmem <- mmapFileByteString geoFile Nothing
either (throwIO . DecodeException) return (openGeoDBBS bsmem)
openGeoDBBS :: BS.ByteString -> Either String GeoDB
openGeoDBBS bsmem = do
hdr <- decode (getHeaderBytes bsmem)
when (hdr ^? key "binary_format_major_version" . geoNum /= (Just 2 :: Maybe Int)) $
Left "Unsupported database version, only v2 supported."
unless (hdr ^? key "record_size" . geoNum `elem` (Just <$> [24, 28, 32 :: Int])) $
Left "Record size not supported."
let res = GeoDB bsmem <$> hdr ^? key "database_type" . _DataString
<*> pure (hdr ^.. key "languages" . _DataArray . traverse . _DataString)
<*> hdr ^? key "node_count" . geoNum
<*> hdr ^? key "record_size" . geoNum
<*> hdr ^? key "ip_version" . toVersion
<*> pure (hdr ^? key "descritpion" . key "en" . _DataString)
maybe (Left "Error decoding header") return res
where
toVersion = geoNum . prism' pfrom pto
where
pfrom :: GeoIP -> Int
pfrom GeoIPv4 = 4
pfrom GeoIPv6 = 6
pto 4 = Just GeoIPv4
pto 6 = Just GeoIPv6
pto _ = Nothing
rawGeoData :: GeoDB -> IP -> Either String GeoField
rawGeoData geodb addr = do
bits <- coerceAddr
offset <- getDataOffset (geoMem geodb, geoDbNodeCount geodb, geoDbRecordSize geodb) bits
strictDataAt offset
where
dataSectionStart = (geoDbRecordSize geodb `div` 4) * fromIntegral (geoDbNodeCount geodb) + 16
dataSection = BS.drop dataSectionStart (geoMem geodb)
strictDataAt :: Int64 -> Either String GeoField
strictDataAt offset = do
raw <- decode (BS.drop (fromIntegral offset) dataSection)
traversePtr (strictDataAt . fromIntegral) raw
coerceAddr
| (IPv4 _) <- addr, GeoIPv4 <- geoDbAddrType geodb = return $ ipToBits addr
| (IPv6 _) <- addr, GeoIPv6 <- geoDbAddrType geodb = return $ ipToBits addr
| (IPv4 addrv4) <- addr, GeoIPv6 <- geoDbAddrType geodb = return $ ipToBits $ IPv6 (ipv4ToIPv6 addrv4)
| otherwise = Left "Cannot search IPv6 address in IPv4 database"
data AS = AS {
asNumber :: Int
, asOrganization :: T.Text
} deriving (Show, Eq)
data GeoResult = GeoResult {
geoContinent :: Maybe T.Text
, geoContinentCode :: Maybe T.Text
, geoCountryISO :: Maybe T.Text
, geoCountry :: Maybe T.Text
, geoLocation :: Maybe Location
, geoCity :: Maybe T.Text
, geoCityConfidence :: Maybe Int
, geoPostalCode :: Maybe T.Text
, geoAS :: Maybe AS
, geoISP :: Maybe T.Text
, geoOrganization :: Maybe T.Text
, geoUserType :: Maybe T.Text
, geoSubdivisions :: [(T.Text, T.Text)]
} deriving (Show, Eq)
data Location = Location {
locationLatitude :: Double
, locationLongitude :: Double
, locationTimezone :: T.Text
, locationAccuracy :: Maybe Int
} deriving (Show, Eq)
findGeoData ::
GeoDB
-> T.Text
-> IP
-> Either String GeoResult
findGeoData geodb lang ip = do
res <- rawGeoData geodb ip
let subdivmap = res ^.. key "subdivisions" . _DataArray . traverse
subdivs = mapMaybe (\s -> (,) <$> s ^? key "iso_code" . _DataString
<*> s ^? key "names" . key lang . _DataString) subdivmap
return $ GeoResult {
geoContinent = res ^? key "continent" . key "names" . key lang . _DataString
, geoContinentCode = res ^? key "continent" . key "code" . _DataString
, geoCountryISO = res ^? key "country" . key "iso_code" . _DataString
, geoCountry = res ^? key "country" . key "names" . key lang . _DataString
, geoLocation = Location <$> res ^? key "location" . key "latitude" . _DataDouble
<*> res ^? key "location" . key "longitude" . _DataDouble
<*> res ^? key "location" . key "time_zone" . _DataString
<*> pure (res ^? key "location" . key "accuracy_radius" . geoNum)
, geoCity = res ^? key "city" . key "names" . key lang . _DataString
, geoCityConfidence = res ^? key "city" . key "confidence" . geoNum
, geoPostalCode = res ^? key "postal" . key "code" . _DataString
, geoAS = AS <$> res ^? key "traits" . key "autonomous_system_number" . geoNum
<*> res ^? key "traits" . key "autonomous_system_organization" . _DataString
, geoISP = res ^? key "traits" . key "isp" . _DataString
, geoOrganization = res ^? key "traits" . key "organization" . _DataString
, geoUserType = res ^? key "traits" . key "user_type" . _DataString
, geoSubdivisions = subdivs
}
key :: T.Text -> Traversal' GeoField GeoField
key k = _DataMap . ix (DataString k)
geoNum :: Num b => Fold GeoField b
geoNum = to fromNum . _Just
where
fromNum (DataInt x) = Just (fromIntegral x)
fromNum (DataWord x) = Just (fromIntegral x)
fromNum _ = Nothing