1 /*
2  * Copyright (C) 2020 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package android.security.cts;
18 
19 
20 import org.junit.runner.Description;
21 import org.junit.runners.model.Statement;
22 
23 import java.util.HashSet;
24 import java.util.Iterator;
25 import java.util.Set;
26 import java.io.File;
27 import java.io.FileNotFoundException;
28 
29 import org.junit.runner.Description;
30 import org.junit.rules.TestWatcher;
31 
32 import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
33 import com.android.tradefed.build.IBuildInfo;
34 import com.android.tradefed.device.DeviceNotAvailableException;
35 import com.android.tradefed.device.ITestDevice;
36 import com.android.tradefed.log.LogUtil.CLog;
37 import com.android.tradefed.testtype.IAbi;
38 
39 import static org.junit.Assume.*;
40 import static org.junit.Assert.*;
41 
42 public class PocPusher extends TestWatcher {
43     private ITestDevice device = null;
44     private CompatibilityBuildHelper buildHelper = null;
45     private IAbi abi = null;
46 
47     private Set<String> filesToCleanup = new HashSet();
48     public boolean bitness32 = true;
49     public boolean bitness64 = true;
50     public boolean appendBitness = true;
51     public boolean cleanup = true;
52 
53     @Override
starting(Description d)54     protected void starting(Description d) {
55         bothBitness();
56         appendBitness = true;
57         cleanup = true;
58     }
59 
60     @Override
finished(Description d)61     protected void finished(Description d) {
62         for (Iterator<String> it = filesToCleanup.iterator(); it.hasNext();) {
63             String file = it.next();
64             try {
65                 CLog.i("Cleaning up %s", file);
66                 device.deleteFile(file);
67             } catch (DeviceNotAvailableException e) {
68                 CLog.e("Device unavailable when cleaning up %s", file);
69                 continue; // try to remove next time
70             }
71             it.remove();
72         }
73     }
74 
setDevice(ITestDevice device)75     public PocPusher setDevice(ITestDevice device) {
76         this.device = device;
77         return this;
78     }
79 
setAbi(IAbi abi)80     public PocPusher setAbi(IAbi abi) {
81         this.abi = abi;
82         return this;
83     }
84 
setBuild(IBuildInfo buildInfo)85     public PocPusher setBuild(IBuildInfo buildInfo) {
86         buildHelper = new CompatibilityBuildHelper(buildInfo);
87         return this;
88     }
89 
appendBitness(boolean append)90     public PocPusher appendBitness(boolean append) {
91         this.appendBitness = append;
92         return this;
93     }
94 
cleanup(boolean cleanup)95     public PocPusher cleanup(boolean cleanup) {
96         this.cleanup = cleanup;
97         return this;
98     }
99 
only32()100     public PocPusher only32() {
101         bitness32 = true;
102         bitness64 = false;
103         return this;
104     }
105 
only64()106     public PocPusher only64() {
107         bitness32 = false;
108         bitness64 = true;
109         return this;
110     }
111 
bothBitness()112     public PocPusher bothBitness() {
113         bitness32 = true;
114         bitness64 = true;
115         return this;
116     }
117 
pushFile(String testFile, String remoteFile)118     public void pushFile(String testFile, String remoteFile)
119             throws FileNotFoundException, DeviceNotAvailableException {
120         if (appendBitness) {
121             // if neither 32 or 64, nothing would ever be pushed.
122             assertTrue("bitness must be 32, 64, or both.", bitness32 || bitness64);
123 
124             String bitness = SecurityTestCase.getAbi(device).getBitness().trim();
125 
126             // 32-bit doesn't have a 64-bit compatibility layer; skipping.
127             assumeFalse(bitness.equals("32") && !bitness32);
128 
129             // push the 32-bit file on 64-bit device if a 64-bit file doesn't exist.
130             if (bitness.equals("64") && !bitness64) {
131                 bitness = "32";
132                 CLog.i("Pushing a 32-bit file onto a 64-bit device.");
133             }
134             testFile += bitness;
135         }
136         CLog.i("Pushing local: %s to remote: %s", testFile.toString(), remoteFile);
137         File localFile = buildHelper.getTestFile(testFile);
138         device.pushFile(localFile, remoteFile);
139         if (cleanup) {
140             filesToCleanup.add(remoteFile);
141         }
142     }
143 }
144