import re
import os.path
import time
from jwcrypto import jwk
from jwcrypto import jwt
from jwcrypto.common import json_encode

grafana_version='9.4.7'
grafana_url='http://localhost:3000/'

# Вытащим приватный ключ из исполняемого файла графаны:
with open('grafana-{}/bin/grafana'.format(grafana_version), mode='rb') as f:
  grafana_bin = f.read()

pat = re.compile(b'.*({"kty":"RSA","kid":"1"[^}]+}).*')

original_pub_start = grafana_bin.find(b'{"kty":"RSA","kid":"1"')
original_pub_end   = grafana_bin.find(b'}', original_pub_start) + 1
original_pub_json  = grafana_bin[original_pub_start:original_pub_end]

# Генерируем новый ключ
new_key_filename = 'new_key.json'
if (not os.path.isfile(new_key_filename)):
  key = jwk.JWK.generate(kty='RSA', size=4096, alg='RS512', kid='1')
  new_pub_json = '{"kty":"RSA","kid":"1","alg":"RS512","n":"' + key.n + '","e":"' + key.e + '"}'
  # Важно что бы длина совпадала с оригиналом
  if len(new_pub_json) == len(original_pub_json):
    new_key_json = key.export()
    with open(new_key_filename, 'w') as text_file:
      text_file.write(new_key_json)
    print('New key generated!')
  else:
    print('Error - key length differs')
else:
  # Загружаем новый ключ из готового файла
  with open(new_key_filename) as f:
    new_key_json = f.read()
  print('New key loaded from file')

key = jwk.JWK.from_json(new_key_json)
new_pub_json = ('{"kty":"RSA","kid":"1","alg":"RS512","n":"' + key.n + '","e":"' + key.e + '"}').encode()
if len(new_pub_json) == len(original_pub_json):
  print('Pub key length Ok!')

# Нужная лицензия
c = {
    'ver': '2.0.0',
    'iss': 'https://grafana.com',
    'sub': grafana_url,
    'san': [],
    'jti': '1337',
    'nbf': int(time.time()),
    'iat': int(time.time()),
    'exp': int(time.time())+10*365*24*60*60,
    'lexp': int(time.time())+10*365*24*60*60,
    'lid': '1337',
    'prod': [
        'grafana-enterprise'
    ],
    'limit_by': 'role_users',
    'included_admins': 100,
    'included_viewers': 1000,
    'included_users': 1000,
    'lic_exp_warn_days': 3,
    'tok_exp_warn_days': 2,
    'update_days': 1,
    'company': 'acme',
    'slug': 'acme',
    'trial': False,
    'usage_billing': False,
    'max_concurrent_user_sessions': 1000,
    'type': 'prod'
}

# Подписываем лицензию новым ключём
Token = jwt.JWT(header={'alg': 'RS512', 'typ': 'JWT', 'kid': '1'},
                 claims=c)
Token.make_signed_token(key)
token_text = Token.serialize()
# Сохраняем лицензию в файл
with open('grafana-{}/bin/license.jwt'.format(grafana_version), 'w') as text_file:
  text_file.write(token_text)
print('License file created')

# Заменяем публичный ключ на новый в исполняемом файле
grafana_bin = grafana_bin.replace(original_pub_json, new_pub_json)
with open('grafana-{}/bin/grafana'.format(grafana_version), mode='wb') as f:
  f.write(grafana_bin)
print('Binary modified')