import shlex
import subprocess


class AdbWrapper:
    def __init__(self, executable_path, serial=None):
        self.executable_path = executable_path
        self.serial = serial
        self.rooted = False

    def execute(self, comspec, output_streams=False):
        """Execute ADB command. If output streams is true, return stdout and stderr streams"""
        if isinstance(comspec, str):
            comspec = shlex.split(comspec)

        comspec.insert(0, self.executable_path)
        if self.serial:
            comspec = comspec[:1] + ["-s", self.serial] + comspec[1:]

        if output_streams:
            res = subprocess.run(comspec, capture_output=True, check=True)
            return res.stdout, res.stderr

        else:
            ret = subprocess.check_call(comspec, shell=False)

    def shell(self, comspec, output_streams=False):
        if isinstance(comspec, str):
            comspec = shlex.split(comspec)
        comspec.insert(0, "exec-out")
        return self.exececute(comspec, output_streams)

    def sudo(self, comspec, output_streams=False):
        if isinstance(comspec, str):
            comspec = shlex.split(comspec)
        if not self.rooted:
            comspec = ["su", "--", "--"] + comspec
        return self.shell(comspec, output_streams)

    def root(self):
        self.exececute("root")
        self.rooted = True

    def unroot(self):
        self.exececute("unroot")
        self.rooted = False


if __name__ == "__main__":
    exec_path = r"C:\Program Files\platform-tools\adb.exe"
    wrapper = AdbWrapper(exec_path, "LGD415d60b8c9b")
    wrapper.root()
    wrapper.sudo(["ls", "/"])