#!/usr/bin/python3

import sys
import os
import logging
import requests
from urllib.parse import urlencode
from io import BytesIO
import ujson
import uuid
import socket
import platform
import time
from cryptography.hazmat.primitives import serialization as crypto_serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.backends import default_backend as crypto_default_backend

logger = logging.getLogger(__name__)


class Authentication(object):
    uuid = None
    apiUrl = None
    clientId = None
    clientSecret = None
    verify = True
    Session = None

    endpointData = None
    tokenData = None

    clientIdFile = '/opt/hawk/etc/hawk-soard-id'

    globalKey = None

    privateKeyFilename = '/opt/hawk/etc/hawk-soard.key'
    privateKey = None

    publicKey = None
    publicKeyFilename = '/opt/hawk/etc/hawk-soard.pub'
    profile = None

    def __init__(self, apiUrl, clientId, clientSecret):
        self.apiUrl = apiUrl
        self.clientId = clientId
        self.clientSecret = clientSecret
        self.Session = requests.Session()
        self.profile = {}

    def setVerify(self, b=True):
        self.verify = self.Session.verify = b

    def getEndpoints(self):
        # get curl info

        url = self.apiUrl + "/api/auth/endpoints/" + self.clientId
        try:
            response = self.Session.get(url, verify=self.verify)
        except Exception as e:
            logger.error("Failed fetching endpoints at url: %s" % url)
            logger.exception(e)
            raise e

        if response.status_code != 200:
            raise Exception("Status response not valid (%d): %s" %
                            (response.status_code, response.content))

        self.endpointData = response.json()

        if not 'token_uri' in self.endpointData:
            data = self.endpointData
            self.Session.close()
            raise Exception(data)

        self.Session.close()
        return True

    def setScope(self, scope=[]):
        if not self.endpointData:
            raise Exception("getEndpoints has not been called yet")
        if not 'scope' in self.endpointData:
            raise Exception("getEndpoints did not return scope")

        self.endpointData['scope'] = scope
        return True

    def getToken(self):

        # get curl info
        # make post and set data request properly
        postData = {"grant_type": "client_credentials", "client_id": self.clientId,
                    "client_secret": self.clientSecret, "scope": " ".join(self.endpointData['scope'])}
        response = self.Session.post(
            self.endpointData['token_uri'], data=postData, verify=self.verify)
        self.tokenData = response.json()
        if not 'access_token' in self.tokenData:
            self.tokenData = None
            self.Session.close()
            raise Exception(
                "No access token found within response: %s" % response.text)
        self.Session.close()
        return True

    def setUUIDFile(self, f):
        self.clientIdFile = f

    def getUUID(self):

        if self.uuid:
            return self.uuid

        try:
            with open(self.clientIdFile) as f:
                self.uuid = f.read()
                f.close()
                return self.uuid
        except:
            pass

        # generate
        self.uuid = str(uuid.uuid4())
        with open(self.clientIdFile, 'w') as f:
            f.write(self.uuid)
            f.close()

        return self.uuid

    def __getDeviceInfo(self):
        # TODO: move to common.common.py from common.common import get_device_info, call hostname, address, platform = get_device_info()
        self.tokenData['hostname'] = socket.gethostname()
        self.tokenData['address'] = socket.gethostbyname(
            self.tokenData['hostname'])
        """
        # new cpe implementation in future
        syst = platform.system().lower()
        if syst == 'windows':
            plat = "microsoft"
        elif syst == "osx":
            plat = "mac"
        elif syst == "linux":
            plat = "linux"
        self.tokenData['platform'] = 'cpe:/o:' + plat.lower() + ':' + platform.platform().replace("-","_") + "::" + platform.version().lower()
        """
        self.tokenData['platform'] = os.name + "/" + \
            platform.system() + " " + platform.release()

    def getParentGroup(self):
        if not isinstance(self.profile, dict):
            raise Exception("Authentication profile is unavailable.")
        if 'group' not in self.profile:
            raise Exception("Authentication profile does not include group.")
        return self.profile['group']

    def loginApiUser(self, access_token, secret_key):
        if not access_token or not secret_key:
            raise Exception("Service-account authentication requires access_token and secret_key.")

        response = self.Session.post(
            self.apiUrl + "/api/auth",
            json={"access_token": access_token, "secret_key": secret_key},
            verify=self.verify,
        )
        try:
            data = response.json()
        except Exception:
            raise Exception("Service-account authentication returned invalid JSON: %s" % response.text)

        if not isinstance(data, dict) or not data.get("status"):
            raise Exception("Service-account authentication failed: %s" % response.text)

        if not isinstance(data.get("data"), dict):
            raise Exception("Service-account authentication missing user profile: %s" % response.text)

        self.profile = dict(data["data"])
        # Mark auth state for runtime paths that check for token existence.
        self.tokenData = {"access_token": "api_token_session"}
        return True

    def registerToken(self):
        if not self.tokenData:
            raise Exception(
                "No token data found, have you called getEndpoints and then getToken?")

        # generate our storage record data
        self.__getDeviceInfo()

        data = self.tokenData

        if not self.publicKey:
            self.loadKeys()

        if not self.publicKey:
            raise Exception("Unable to load public key")

        # TODO: Actual calls to the api should all be in one file in soar.py that way we can just have everything in one spot.
        data['publicKey'] = self.publicKey

        response = self.Session.post(self.apiUrl +
                                     "/api/auth/token/" + self.clientId + "/" + self.getUUID(),
                                     data=self.tokenData, verify=self.verify)
        try:
            results = response.json()
        except Exception as e:
            self.Session.close()
            raise e

        if not results['status']:
            raise Exception(results['details'])

        self.profile = results['data']

        if not results['data']['approval']:
            self.Session.close()
            return False

        return True

    def setPrivateKey(self, fn):
        self.privateKeyFilename = fn

    def setPublicKey(self, fn):
        self.publicKeyFilename = fn

    def loadKeys(self):
        self.loadPrivateKey()
        self.loadPublicKey()

    def loadPrivateKey(self):
        self.privateKey = None
        self.globalKey = None
        if not os.path.isfile(self.privateKeyFilename):
            self.generateKeys()
            self.saveKeys()
            return True

        with open(self.privateKeyFilename, 'rb') as pem_in:
            pemlines = pem_in.read()
            self.globalKey = crypto_serialization.load_pem_private_key(
                pemlines, None, crypto_default_backend())

            self.privateKey = self.globalKey.private_bytes(
                crypto_serialization.Encoding.PEM,
                crypto_serialization.PrivateFormat.PKCS8,
                crypto_serialization.NoEncryption())

    def loadPublicKey(self):
        self.publicKey = None
        if not os.path.isfile(self.publicKeyFilename):
            if not self.globalKey:
                raise Exception(
                    "No RSA private key available to export public key")

            self.publicKey = self.globalKey.public_key().public_bytes(
                crypto_serialization.Encoding.OpenSSH,
                crypto_serialization.PublicFormat.OpenSSH)
            self.savePublicKey()

        if not os.path.isfile(self.publicKeyFilename):
            raise Exception("Unable to save public key to disk: %s" %
                            self.publicKeyFilename)

        with open(self.publicKeyFilename, 'rb') as pem_in:
            pemlines = pem_in.read()
            self.publicKey = pemlines

        return True

    def generateKeys(self):

        if self.privateKey and self.publicKey:
            return True

        self.globalKey = rsa.generate_private_key(
            backend=crypto_default_backend(),
            public_exponent=65537,
            key_size=2048)

        self.privateKey = self.globalKey.private_bytes(
            crypto_serialization.Encoding.PEM,
            crypto_serialization.PrivateFormat.PKCS8,
            crypto_serialization.NoEncryption())

        self.publicKey = self.globalKey.public_key().public_bytes(
            crypto_serialization.Encoding.OpenSSH,
            crypto_serialization.PublicFormat.OpenSSH)

        return True

    def saveKeys(self):
        self.savePublicKey()
        self.savePrivateKey()

    def savePrivateKey(self):
        pem = self.globalKey.private_bytes(
            encoding=crypto_serialization.Encoding.PEM,
            format=crypto_serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=crypto_serialization.NoEncryption())
        with open(self.privateKeyFilename, 'wb') as pem_out:
            pem_out.write(pem)

    def savePublicKey(self):
        with open(self.publicKeyFilename, 'wb') as pem_out:
            pem_out.write(self.publicKey)

    def finish(self):
        return self.__finished()

    def __finished(self):
        try:
            self.Session.close()
        except:
            pass

        return True


if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage %s: <host> <clientId> <secret>" % sys.argv[0])
        sys.exit(-1)

    a = Authentication(sys.argv[1], sys.argv[2], sys.argv[3])

    a.setVerify(False)

    a.getEndpoints()

    a.setScope(["https://graph.microsoft.com/.default"])
    if a.getToken():
        print(a.tokenData)

        while not a.registerToken():
            time.sleep(30)

        a.finish()
