From e57f08edbbd91dff7d03bc42f67135dc83a3cb50 Mon Sep 17 00:00:00 2001
From: Difrex <difrex.punk@gmail.com>
Date: Tue, 11 Oct 2016 14:36:41 +0300
Subject: [PATCH] Closed SD-10

---
 surok/discovery.py | 65 ++++++++++++++++++++++++++++++++--------------
 1 file changed, 46 insertions(+), 19 deletions(-)

diff --git a/surok/discovery.py b/surok/discovery.py
index 8042a1f..604c441 100644
--- a/surok/discovery.py
+++ b/surok/discovery.py
@@ -2,12 +2,18 @@ import dns.resolver
 import dns.query
 from dns.exception import DNSException
 import logging
+import sys
 
 
-# Configure logging
-FORMAT = '%(asctime) %(message)s'
-logging.basicConfig(format=FORMAT)
-logger = logging.getLogger('discovery')
+# Logger configuration
+# This need to be moved
+def get_logger():
+    # Configure logging
+    FORMAT = '%(asctime) %(message)s'
+    logging.basicConfig(format=FORMAT)
+    logger = logging.getLogger(__name__)
+    return logger
+
 
 # Resolve service from mesos-dns SRV record
 # return dict {"servicename": [{"name": "service.f.q.d.n.", "port": 9999}]}
@@ -15,23 +21,17 @@ def resolve(app, conf):
     hosts = {}
     services = app['services']
     domain = conf['domain']
-    group = None
-
-    # Groups hack
-    if app['env'].get('SUROK_DISCOVERY_GROUP') is not None:
-        group = app['env']['SUROK_DISCOVERY_GROUP']
+    logger = get_logger()
 
     for service in services:
         hosts[service['name']] = {}
 
-        # Check group configuration
-        if group is not None:
-            pass
-        else:
-            # Load group from service config
-            # /etc/surok/conf.d/service_conf.json
-            group = service['group']
-
+        group = get_group(service, app)
+        if group is False:
+            logger.error('Group is not defined in config, SUROK_DISCOVERY_GROUP and MARATHON_APP_ID')
+            logger.error('Not in Mesos launch?')
+            sys.exit(2)
+        
         # Port name from app config
         ports = None
         try:
@@ -56,6 +56,7 @@ def resolve(app, conf):
 # Do SRV queries
 # Return array: [{"name": "f.q.d.n", "port": 8876}]
 def do_query(fqdn, loglevel):
+    logger = get_logger()
     servers = []
     try:
         query = dns.resolver.query(fqdn, 'SRV')
@@ -75,5 +76,31 @@ def do_query(fqdn, loglevel):
 
 # Groups switch
 # Priority: config, environment, marathon environment
-def get_group(conf):
-    pass
+def get_group(service, app):
+    # Check group in app conf
+    if 'group' in service:
+        return service['group']
+    # Check environment variable
+    elif app['env'].get('SUROK_DISCOVERY_GROUP'):
+        return app['env']['SUROK_DISCOVERY_GROUP']
+    # Check marathon environment variable
+    elif app['env'].get('MARATHON_APP_ID'):
+        group = parse_marathon_app_id(app['env']['MARATHON_APP_ID'])
+        return group
+    else:
+        return False
+
+
+# Parse MARATHON_APP_ID
+def parse_marathon_app_id(marathon_app_id):
+    marathon_app_id = marathon_app_id.split('/')
+    group = ''
+    counter = len(marathon_app_id) - 2
+    i = 0
+    while counter > i:
+        group = group + marathon_app_id[counter]
+        if counter != i + 1:
+            group += '.'
+        counter -= 1
+
+    return group