####
# detect the leaked gcp disk on dev environments
# it will be added into gcloud function later
###
import json
import subprocess
import os
import argparse
import yaml
from kubernetes import client, config
import base64
import time
from hashlib import sha256
import requests as httpRequests
import hmac
import csv
from datetime import datetime
import googleapiclient
from googleapiclient import discovery
from tempfile import TemporaryFile
import google.auth
from google.auth.transport import requests
from oauth2client.client import GoogleCredentials

gkeClusters = {
  "staging-seed-asia-northeast1-83d8caa":  "asia-northeast1",
  "staging-seed-asia-southeast1-b7bc5b9": "asia-southeast1",
  "staging-seed-us-central1-6159679":  "us-central1",
  "dev-seed-us-west1-3137aff":    "us-west1",
  "stability-seed-us-west1-67ac24a":  "us-west1",
  "staging-seed-us-west1-73cb72c":     "us-west1",
  "dev-seed-us-west2-9f7b847":     "us-west2"
}

gkeAddresses = {
  "staging-seed-asia-northeast1-83d8caa":  "35.190.235.134",
  "staging-seed-asia-southeast1-b7bc5b9": "35.240.161.172",
  "staging-seed-us-central1-6159679":  "35.223.148.203",
  "dev-seed-us-west1-3137aff":    "34.82.26.237",
  "stability-seed-us-west1-67ac24a":  "35.227.149.173",
  "staging-seed-us-west1-73cb72c":     "35.227.149.173",
  "dev-seed-us-west2-9f7b847":     "35.236.101.158"
}

outfile='disks-{}.csv'.format(datetime.now().strftime('%Y-%m-%d-%H:%M:%S') )
projectName = 'pingcap-gardener'
bucketName = 'daily-resource-leak-detetor'
allPVs = []
disksDict = {}
leakedNumer = 0
# feishu_webhook = 'https://open.feishu.cn/open-apis/bot/v2/hook/402619ca-3aad-4d9f-86a4-8e1bea90bb76'
feishu_webhook = os.getenv("feishu_webhook")
feishu_secret = os.getenv("feishu_secret")

csvFile = open( '/tmp/' + outfile, 'w', newline='')
fieldnames = ['id', 'name', 'zone', 'selfLink', "sizeGb", "creationTimestamp", 'lastAttachTimestamp']
writer = csv.DictWriter(csvFile, fieldnames=fieldnames)
writer.writeheader()

def get_disks():
  computeClient = googleapiclient.discovery.build('compute', 'v1')
  credentials = GoogleCredentials.get_application_default()
  disksRequest = computeClient.disks().aggregatedList(project="pingcap-gardener")
  while disksRequest is not None:
    diskResponse = disksRequest.execute()
    for name, disks_scoped_list in diskResponse['items'].items():
      if 'disks' in disks_scoped_list and len(disks_scoped_list['disks']) > 0:
        for k, v in disks_scoped_list['disks'][0].items():
          disksDict[k] = v
    disksRequest = computeClient.disks().aggregatedList_next(previous_request=disksRequest, previous_response=diskResponse)

def update_gke_kubeconfig(gkeName, gkeZone):
  credentials = GoogleCredentials.get_application_default()
  containerClient = googleapiclient.discovery.build('container', 'v1', credentials=credentials)
  clusterName = "projects/{}/locations/{}/clusters/{}".format(projectName, gkeZone, gkeName)

  gke_clusters = containerClient.projects().locations().clusters()
  gke_cluster = gke_clusters.get(name=clusterName).execute()


def read_all_persistent_volume(k8sIp):
  credentials, project_id = google.auth.default()
  session = requests.AuthorizedSession(credentials)
  response = session.get('https://{}/api/v1/persistentvolumes'.format(k8sIp), verify=False)
  response.raise_for_status()
  return response.json()

def process_gke_clusters():
  for gke in gkeClusters:
    pvs = read_all_persistent_volume(gkeAddresses[gke])
    for pv in pvs['items']:
      print("process gke {} pv {}".format(gke, pv['metadata']['name']))
      allPVs.append(pv['metadata']['name'])

def read_kubeconfig_secret(k8sIp, ns, secret):
  credentials, project_id = google.auth.default()
  session = requests.AuthorizedSession(credentials)
  response = session.get('https://{}/api/v1/namespaces/{}/secrets/{}'.format(k8sIp, ns, secret), verify=False)
  response.raise_for_status()
  return response.json()

def read_namespaces(k8sIp):
  credentials, project_id = google.auth.default()
  session = requests.AuthorizedSession(credentials)
  response = session.get('https://{}/api/v1/namespaces'.format(k8sIp), verify=False)
  response.raise_for_status()
  return response.json()

def process_all_shoot_clusters():
  for gke in gkeClusters:
    # update_gke_kubeconfig(gke, gkeClusters[gke])
    # config.kube_config.load_kube_config()
    # apiClient = client.CoreV1Api()
    #
    # namespaces = apiClient.list_namespace()
    namespaces = read_namespaces(gkeAddresses[gke])
    print(namespaces)

    for ns in namespaces['items']:
      # update_gke_kubeconfig(gke, gkeClusters[gke])
      # config.kube_config.load_kube_config()
      # apiClient = client.CoreV1Api()

      if 'shoot--' not in ns['metadata']['name']:
        continue
      print("process shoot {}".format( ns['metadata']['name'] ))
      try:
        secret = read_kubeconfig_secret(gkeAddresses[gke], ns, 'kubecfg')
        print(secret)
        shootKubeconfig = apiClient.read_namespaced_secret('kubecfg', ns['metadata']['name'])
        shootKubeconfig = shootKubeconfig.data["kubeconfig"]
        shootKubeconfig= base64.b64decode(shootKubeconfig)
        shootKubeconfigFile= open("/tmp/" + ns['metadata']['name'] , "w")
        yaml.dump(yaml.load(shootKubeconfig), shootKubeconfigFile)
        shootKubeconfigFile.close()
        config.kube_config.load_kube_config(config_file= "/tmp/" + ns['metadata']['name'])
        apiClient = client.CoreV1Api()
        pvs = apiClient.list_persistent_volume()
        for pv in pvs.items:
          print("process shoot {} pv {}".format(ns['metadata']['name'], pv.metadata.name))
          allPVs.append(pv.metadata.name)
      except:
        print("bypass shoot {}".format(ns['metadata']['name']))

def process():
  for instance in allPVs:
    print(instance)

  for disk in disksDict:
    try:
      foundIt = False
      targetPVName = ""
      if 'description' in disk:
        description = json.loads(disk['description'])
        if 'kubernetes.io/created-for/pv/name' in description:
          targetPVName = description['kubernetes.io/created-for/pv/name']

      if targetPVName != "":
        for pv in allPVs:
          if pv == targetPVName:
            foundIt = True

      if targetPVName != "" and not foundIt:
        writer.writerow({
          'id': disk['id'],
          'name': disk['name'],
          'zone': disk['zone'],
          'selfLink': disk['selfLink'],
          "sizeGb": disk["sizeGb"],
          "creationTimestamp": disk["creationTimestamp"],
          'lastAttachTimestamp': disk['lastAttachTimestamp']})
        leakedNumer = leakedNumer + 1
    except:
      pass

  csvFile.close()

def get_sign_timestamp():
  timestamp = str(int(time.time()))
  # secret = os.getenv("feishu_secret")

  key = f'{timestamp}\n{feishu_secret}'
  key_enc = key.encode('utf-8')
  msg = ""
  msg_enc = msg.encode('utf-8')
  hmac_code = hmac.new(key_enc, msg_enc, digestmod=sha256).digest()
  sign = base64.b64encode(hmac_code).decode('utf-8')
  return {'timestamp': timestamp, 'sign': sign}


def post_bot_webhook():
  sign_timestamp = get_sign_timestamp()
  payload_message = {
    "timestamp": sign_timestamp['timestamp'],
    "sign": sign_timestamp['sign'],
    "msg_type": "interactive",
    "card": {
      "i18n_elements": {
        "en_us": [
          {
            "tag": "div",
            "text": {
              "tag": "lark_md",
              "content": "**可能泄露的GCP硬盘有{}块**".format(leakedNumer)
            },
          },
          {
            "tag": "div",
            "text": {
              "tag": "lark_md",
              "content": "**结果存放在gcp对象存储：{}/{} **".format(bucketName, outfile)
            },
          },
        ]
      }
    }
  }
  headers = {
    'Content-Type': 'application/json'
  }

  response = httpRequests.request(
    "POST", feishu_webhook, headers=headers, data=json.dumps(payload_message))
  print(response.text)


def upload_result_to_gcp_storage_and_send_alerts():
  gcpBlobClient = googleapiclient.discovery.build('storage', 'v1')

  body = {
    'name': outfile,
  }

  with open('/tmp/' + outfile, 'rb') as f:
    req = gcpBlobClient.objects().insert(
      bucket='daily-resource-leak-detetor', body=body,
      media_body=googleapiclient.http.MediaIoBaseUpload(
        f, 'application/octet-stream'))
    resp = req.execute()

  post_bot_webhook()

get_disks()
process_gke_clusters()
process_all_shoot_clusters()
process()
upload_result_to_gcp_storage_and_send_alerts()