import os

from keystoneclient import session as ksc_session
from keystoneclient.auth.identity import v3
from keystoneclient.v3 import client as keystone_v3

try:
    # Used for creating the ADMIN user
    OS_PASSWORD = '123456'
    OS_USERNAME = 'admin'
    # This will vary according to the entity:
    # the IdP or the SP
    OS_AUTH_URL = 'http://sp:35357/v3'
    OS_PROJECT_NAME = 'admin'
    OS_DOMAIN_NAME = 'default'
except KeyError as e:
    raise SystemExit('%s environment variable not set.' % e)

def client_for_admin_user():
    auth = v3.Password(auth_url=OS_AUTH_URL,
                       username=OS_USERNAME,
                       password=OS_PASSWORD,
                       user_domain_name=OS_DOMAIN_NAME,
                       project_name=OS_PROJECT_NAME,
                       project_domain_name=OS_DOMAIN_NAME)
    session = ksc_session.Session(auth=auth)
    return keystone_v3.Client(session=session)

# Used to execute all admin actions
client = client_for_admin_user()

def create_domain(client, name):
    try:
         d = client.domains.create(name=name)
    except:
         d = client.domains.find(name=name)
    return d


def create_group(client, name, domain):
    try:
         g = client.groups.create(name=name, domain=domain)
    except:
         g = client.groups.find(name=name)
    return g


def create_role(client, name):
    try:
        r = client.roles.create(name=name)
    except:
        r = client.roles.find(name=name)
    return r


print('\nCreating domain1')
domain1 = create_domain(client, 'domain1')

print('\nCreating group1')
group1 = create_group(client, 'group1', domain1)

print('\nCreating role Member')
role1 = create_role(client, 'Member')

print('\nGrant role Member to group1 in domain1')
client.roles.grant(role1, group=group1, domain=domain1)

print('\nList group1 role assignments')
client.role_assignments.list(group=group1)

def create_mapping(client, mapping_id, rules):
    try:
        m = client.federation.mappings.create(
            mapping_id=mapping_id, rules=rules)
    except:
        m = client.federation.mappings.find(
            mapping_id=mapping_id)
    return m

print('\nCreating mapping')
rules = [
{
    "local": [
        {
            "user": {
                "name": "federated_user"
            },
            "group": {
                "id": group1.id
            }
        }
    ],
    "remote": [
        {
            "type": "openstack_user",
            "any_one_of": [
                "bob",
                "acme"
            ]
        }
    ]
}
]

mapping1 = create_mapping(client, mapping_id='keystone-idp-mapping', rules=rules)


def create_idp(client, id, remote_id):
    idp_ref = {'id': id,
               'remote_ids': [remote_id],
               'enabled': True}
    try:
        i = client.federation.identity_providers.create(**idp_ref)
    except:
        i = client.federation.identity_providers.find(id=id)
    return i

def create_protocol(client, protocol_id, idp, mapping):
    try:
        p = client.federation.protocols.create(protocol_id=protocol_id,
                                               identity_provider=idp,
                                               mapping=mapping)
    except:
        p = client.federation.protocols.find(protocol_id=protocol_id)
    return p


print('\nRegister keystone-idp')
idp1 = create_idp(client, id='keystone-idp',
                  remote_id='http://idp:5000/v3/OS-FEDERATION/saml2/idp')

print('\nRegister protocol')
protocol1 = create_protocol(client, protocol_id='saml2', idp=idp1,
                            mapping=mapping1)

import json
import os

import requests

from keystoneclient import session as ksc_session
from keystoneclient.auth.identity import v3
from keystoneclient.v3 import client as keystone_v3


class K2KClient(object):
    def __init__(self):
        self.sp_id = 'keystone-sp'
        self.auth_url = 'http://idp:35357/v3'
        self.project_id = '582df27a0db14149a6da375b31fce3df'
        self.username = 'bob'
        self.password = '123456'
        self.domain_id = 'default'

    def v3_authenticate(self):
        auth = v3.Password(auth_url=self.auth_url,
                           username=self.username,
                           password=self.password,
                           user_domain_id=self.domain_id,
                           project_id=self.project_id)
        self.session = ksc_session.Session(auth=auth, verify=False)
        self.session.auth.get_auth_ref(self.session)
        self.token = self.session.auth.get_token(self.session)
    def _generate_token_json(self):
        return {
            "auth": {
                "identity": {
                    "methods": [
                        "token"
                    ],
                    "token": {
                        "id": self.token
                    }
                },
                "scope": {
                    "service_provider": {
                        "id": self.sp_id
                    }
                }
            }
        }

    def _check_response(self, response):
        if not response.ok:
            raise Exception("Something went wrong, %s" % response.__dict__)

    def get_saml2_ecp_assertion(self):
        token = json.dumps(self._generate_token_json())
        url = self.auth_url + '/auth/OS-FEDERATION/saml2/ecp'
        r = self.session.post(url=url, data=token, verify=False)
        self._check_response(r)
        self.assertion = str(r.text)

    def _get_sp(self):
        url = self.auth_url + '/OS-FEDERATION/service_providers/' + self.sp_id
        r = self.session.get(url=url, verify=False)
        self._check_response(r)
        sp = json.loads(r.text)[u'service_provider']
        return sp

    def _handle_http_302_ecp_redirect(self, response, location, **kwargs):
        return self.session.get(location, authenticated=False, **kwargs)

    def exchange_assertion(self):
        """Send assertion to a Keystone SP and get token."""
        sp = self._get_sp()

        r = self.session.post(
            sp[u'sp_url'],
            headers={'Content-Type': 'application/vnd.paos+xml'},
            data=self.assertion,
            authenticated=False,
            redirect=False)

        self._check_response(r)

        r = self._handle_http_302_ecp_redirect(r, sp[u'auth_url'],
                                               headers={'Content-Type':
                                               'application/vnd.paos+xml'})
        self.fed_token_id = r.headers['X-Subject-Token']
        self.fed_token = r.text

    def list_federated_projects(self):
        url = 'http://sp:5000/v3/OS-FEDERATION/projects'
        headers = {'X-Auth-Token': self.fed_token_id}
        r = requests.get(url=url, headers=headers)
        self._check_response(r)
        return json.loads(str(r.text))

    def _get_scoped_token_json(self, project_id):
        return {
            "auth": {
                "identity": {
                    "methods": [
                        "token"
                    ],
                    "token": {
                        "id": self.fed_token_id
                    }
                },
                "scope": {
                    "project": {
                        "id": project_id
                    }
                }
            }
        }

    def scope_token(self, project_id):
        # project_id can be select from the list in the previous step
        token = json.dumps(self._get_scoped_token_json(project_id))
        url = 'http://sp:5000/v3/auth/tokens'
        headers = {'X-Auth-Token': self.fed_token_id,
                   'Content-Type': 'application/json'}
        r = requests.post(url=url, headers=headers, data=token,
                          verify=False)
        self._check_response(r)
        self.scoped_token_id = r.headers['X-Subject-Token']
        self.scoped_token = str(r.text)


def main():
    client = K2KClient()
    client.v3_authenticate()
    client.get_saml2_ecp_assertion()
    print('ECP wrapped SAML assertion: %s' % client.assertion)
    client.exchange_assertion()
    print('Unscoped token id: %s' % client.fed_token_id)
    
    # If you want to get a scope token, please ensure federated_user has a project
    # and uncommen below codes.
    '''
    projects = client.list_federated_projects()
    print('Federated projects: %s' % projects['projects'])
    project_id = projects['projects'][0]['id']
    project_name = projects['projects'][0]['name']
    client.scope_token(project_id)
    print('Scoped token of ' + project_name + ' : ' + client.scoped_token_id)
    '''


if __name__ == "__main__":
    main()