/*
 * Copyright (C) 2017 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 */

package android.jvmti.cts;

import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
import com.android.ddmlib.testrunner.RemoteAndroidTestRunner;
import com.android.tradefed.build.IBuildInfo;
import com.android.tradefed.config.Option;
import com.android.tradefed.device.ITestDevice;
import com.android.tradefed.log.LogUtil.CLog;
import com.android.tradefed.result.ITestLifeCycleReceiver;
import com.android.tradefed.result.TestDescription;
import com.android.tradefed.testtype.DeviceTestCase;
import com.android.tradefed.testtype.IAbi;
import com.android.tradefed.testtype.IAbiReceiver;
import com.android.tradefed.testtype.IBuildReceiver;
import com.android.tradefed.util.AbiUtils;
import com.android.tradefed.util.FileUtil;
import com.android.tradefed.util.ZipUtil;

import java.io.File;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import java.util.zip.ZipFile;

/**
 * Test a JVMTI device test.
 *
 * Reads the configuration (APK and package name) out of the embedded config.properties file. Runs
 * the agent (expected to be packaged with the APK) into the app's /data/data directory, starts a
 * test run and attaches the agent.
 */
public class JvmtiHostTest extends DeviceTestCase implements IBuildReceiver, IAbiReceiver {
    private static final String RUNNER = "androidx.test.runner.AndroidJUnitRunner";
    // inject these options from HostTest directly using --set-option <option name>:<option value>
    @Option(name = "package-name",
            description = "The package name of the device test",
            mandatory = true)
    private String mTestPackageName = null;

    @Option(name = "test-file-name",
            description = "the name of a test zip file to install on device.",
            mandatory = true)
    private String mTestApk = null;

    @Option(name = "hidden-api-checks",
            description = "If we should enable hidden api checks. Default 'true'. Set to 'false' " +
            "to disable hiddenapi.",
            mandatory = false)
    private String mHiddenApiChecksEnabled = null;

    private CompatibilityBuildHelper mBuildHelper;
    private IAbi mAbi;
    private int mCurrentUser;

    @Override
    public void setBuild(IBuildInfo arg0) {
        mBuildHelper = new CompatibilityBuildHelper(arg0);
    }

    @Override
    public void setAbi(IAbi arg0) {
        mAbi = arg0;
    }

    // Constant returned to indicate get-current-user failed. See comment at/near
    // https://cs.android.com/android/_/android/platform/tools/tradefederation/+/android11-release:device_build_interfaces/com/android/tradefed/device/ITestDevice.java;l=780
    private static final int GET_USER_FAILURE = -10000;

    // Try getting current user and throw an exception immediately if we fail.
    @Override
    protected void setUp() throws Exception {
        mCurrentUser = getDevice().getCurrentUser();
        if (mCurrentUser == GET_USER_FAILURE) {
            throw new RuntimeException("am get-current-user failed!");
        }
    }

    public void testJvmti() throws Exception {
        final ITestDevice device = getDevice();

        String testingArch = AbiUtils.getBaseArchForAbi(mAbi.getName());
        String deviceArch = getDeviceBaseArch(device);

        //Only bypass if Base Archs are different
        if (!testingArch.equals(deviceArch)) {
            CLog.d(
                    "Bypass as testing Base Arch:"
                            + testingArch
                            + " is different from DUT Base Arch:"
                            + deviceArch);
            return;
        }

        if (mTestApk == null || mTestPackageName == null) {
            throw new IllegalStateException("Incorrect configuration");
        }

        if (null != mHiddenApiChecksEnabled &&
            !"false".equals(mHiddenApiChecksEnabled) &&
            !"true".equals(mHiddenApiChecksEnabled)) {
          throw new IllegalStateException(
              "option hidden-api-checks must be 'true' or 'false' if present.");
        }
        boolean disable_hidden_api =
            mHiddenApiChecksEnabled != null && "false".equals(mHiddenApiChecksEnabled);
        String old_hiddenapi_setting = null;
        if (disable_hidden_api) {
            old_hiddenapi_setting = device.getSetting("global", "hidden_api_policy");
            device.setSetting("global", "hidden_api_policy", "1");
        }

        try {
            RemoteAndroidTestRunner runner = new RemoteAndroidTestRunner(mTestPackageName, RUNNER,
                    device.getIDevice());
            // set a max deadline limit to avoid hanging forever
            runner.setMaxTimeToOutputResponse(10, TimeUnit.MINUTES);

            AttachAgent aa = new AttachAgent(device, mTestPackageName, mTestApk);
            aa.prepare();
            TestResults tr = new TestResults(aa);

            device.runInstrumentationTests(runner, tr);

            assertTrue(tr.getErrors(), tr.hasStarted());
            assertFalse(tr.getErrors(), tr.hasFailed());
        } finally {
            if (disable_hidden_api) {
                device.setSetting("global", "hidden_api_policy", old_hiddenapi_setting);
            }
        }
    }

    private String getDeviceBaseArch(ITestDevice device) throws Exception {
        String abi = device.executeShellCommand("getprop ro.product.cpu.abi").replace("\n", "");
        CLog.d("DUT abi:" + abi);
        return AbiUtils.getBaseArchForAbi(abi);
    }

    private class AttachAgent implements Runnable {
        private ITestDevice mDevice;
        private String mPkg;
        private String mApk;

        private String mAgentInDataData;

        public AttachAgent(ITestDevice device, String pkg, String apk) {
            this.mDevice = device;
            this.mPkg = pkg;
            this.mApk = apk;
        }

        public void prepare() {
            try {
                String pwd = mDevice.executeShellCommand(
                        "run-as " + mPkg + " --user " + mCurrentUser + " pwd");
                if (pwd == null) {
                    throw new RuntimeException("pwd failed");
                }
                pwd = pwd.trim();
                if (pwd.isEmpty()) {
                    throw new RuntimeException("pwd failed");
                }

                mAgentInDataData = installLibToDataData(pwd, "libctsjvmtiagent.so");
            } catch (Exception e) {
                throw new RuntimeException("Failed installing", e);
            }
        }

        @Override
        public void run() {
            try {
                if (mAgentInDataData == null) {
                    throw new IllegalStateException("prepare() has not been called");
                }
                String attachCmd = "cmd activity attach-agent " + mPkg + " " + mAgentInDataData;
                String attachReply = mDevice.executeShellCommand(attachCmd);
                // Don't try to parse the output. The test will time out anyways if this didn't
                // work.
                if (attachReply != null && !attachReply.trim().isEmpty()) {
                    CLog.e(attachReply);
                }
            } catch (Exception e) {
                throw new RuntimeException("Failed attaching", e);
            }
        }

        private String installLibToDataData(String dataData, String library) throws Exception {
            ZipFile zf = null;
            File tmpFile = null;
            String libInTmp = null;
            try {
                String libInDataData = dataData + "/" + library;

                File apkFile = mBuildHelper.getTestFile(mApk);
                zf = new ZipFile(apkFile);

                String libPathInApk = "lib/" + mAbi.getName() + "/" + library;
                tmpFile = ZipUtil.extractFileFromZip(zf, libPathInApk);

                libInTmp = "/data/local/tmp/" + tmpFile.getName();
                if (!mDevice.pushFile(tmpFile, libInTmp)) {
                    throw new RuntimeException("Could not push library " + library + " to device");
                }

                String runAsCp = mDevice.executeShellCommand(
                        "run-as " + mPkg + " --user " + mCurrentUser +
                                " cp " + libInTmp + " " + libInDataData);
                if (runAsCp != null && !runAsCp.trim().isEmpty()) {
                    throw new RuntimeException(runAsCp.trim());
                }

                String runAsChmod = mDevice.executeShellCommand(
                        "run-as " + mPkg + " --user " + mCurrentUser +
                                " chmod a+x " + libInDataData);
                if (runAsChmod != null && !runAsChmod.trim().isEmpty()) {
                    throw new RuntimeException(runAsChmod.trim());
                }

                return libInDataData;
            } finally {
                FileUtil.deleteFile(tmpFile);
                ZipUtil.closeZip(zf);
                if (libInTmp != null) {
                    try {
                        mDevice.executeShellCommand("rm " + libInTmp);
                    } catch (Exception e) {
                        CLog.e("Failed cleaning up library on device");
                    }
                }
            }
        }
    }

    private static class TestResults implements ITestLifeCycleReceiver {
        private boolean mFailed = false;
        private boolean mStarted = false;
        private final Runnable mOnStart;
        private List<String> mErrors = new LinkedList<>();

        public TestResults(Runnable onStart) {
            this.mOnStart = onStart;
        }

        public boolean hasFailed() {
            return mFailed;
        }

        public boolean hasStarted() {
            return mStarted;
        }

        public String getErrors() {
            if (mErrors.isEmpty()) {
                return "";
            }
            return mErrors.toString();
        }

        @Override
        public void testAssumptionFailure(TestDescription arg0, String arg1) {
            mFailed = true;
            mErrors.add(arg0.toString() + " " + arg1);
        }

        @Override
        public void testEnded(TestDescription arg0, Map<String, String> arg1) {}

        @Override
        public void testFailed(TestDescription arg0, String arg1) {
            mFailed = true;
            mErrors.add(arg0.toString() + " " + arg1);
        }

        @Override
        public void testIgnored(TestDescription arg0) {}

        @Override
        public void testRunEnded(long arg0, Map<String, String> arg1) {}

        @Override
        public void testRunFailed(String arg0) {
            mFailed = true;
            mErrors.add(arg0);
        }

        @Override
        public void testRunStarted(String arg0, int arg1) {
            if (mOnStart != null) {
                mOnStart.run();
            }
        }

        @Override
        public void testRunStopped(long arg0) {}

        @Override
        public void testStarted(TestDescription arg0) {
            mStarted = true;
        }
    }
}