James Moger
2015-11-22 ed552ba47c02779c270ffd62841d6d1048dade70
commit | author | age
521cb6 1 /*
JM 2  * Copyright 2014 gitblit.com.
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 package com.gitblit.tests;
17
18 import java.io.ByteArrayInputStream;
19 import java.io.ByteArrayOutputStream;
20 import java.io.IOException;
21 import java.io.OutputStreamWriter;
22 import java.io.Writer;
23 import java.net.SocketAddress;
24 import java.security.KeyPair;
25 import java.security.KeyPairGenerator;
26 import java.security.PublicKey;
27 import java.util.concurrent.atomic.AtomicBoolean;
28
29 import org.apache.sshd.client.ServerKeyVerifier;
d41034 30 import org.apache.sshd.client.SshClient;
JM 31 import org.apache.sshd.client.channel.ClientChannel;
32 import org.apache.sshd.client.session.ClientSession;
521cb6 33 import org.apache.sshd.common.util.SecurityUtils;
JM 34 import org.junit.After;
35 import org.junit.AfterClass;
36 import org.junit.Before;
37 import org.junit.BeforeClass;
38
39 import com.gitblit.Constants.AccessPermission;
40 import com.gitblit.transport.ssh.IPublicKeyManager;
41 import com.gitblit.transport.ssh.MemoryKeyManager;
42 import com.gitblit.transport.ssh.SshKey;
43
44 /**
45  * Base class for SSH unit tests.
46  */
47 public abstract class SshUnitTest extends GitblitUnitTest {
48
49     protected static final AtomicBoolean started = new AtomicBoolean(false);
50     protected static KeyPairGenerator generator;
51     protected KeyPair rwKeyPair;
52     protected KeyPair roKeyPair;
53     protected String username = "admin";
54     protected String password = "admin";
55
56     @BeforeClass
57     public static void startGitblit() throws Exception {
58         generator = SecurityUtils.getKeyPairGenerator("RSA");
59         started.set(GitBlitSuite.startGitblit());
60     }
61
62     @AfterClass
63     public static void stopGitblit() throws Exception {
64         if (started.get()) {
65             GitBlitSuite.stopGitblit();
66         }
67     }
68
69     protected MemoryKeyManager getKeyManager() {
70         IPublicKeyManager mgr = gitblit().getPublicKeyManager();
71         if (mgr instanceof MemoryKeyManager) {
72             return (MemoryKeyManager) gitblit().getPublicKeyManager();
73         } else {
74             throw new RuntimeException("unexpected key manager type " + mgr.getClass().getName());
75         }
76     }
77
78     @Before
79     public void prepare() {
80         rwKeyPair = generator.generateKeyPair();
81
82         MemoryKeyManager keyMgr = getKeyManager();
83         keyMgr.addKey(username, new SshKey(rwKeyPair.getPublic()));
84
85         roKeyPair = generator.generateKeyPair();
86         SshKey sshKey = new SshKey(roKeyPair.getPublic());
87         sshKey.setPermission(AccessPermission.CLONE);
88         keyMgr.addKey(username, sshKey);
89     }
90
91     @After
92     public void tearDown() {
93         MemoryKeyManager keyMgr = getKeyManager();
94         keyMgr.removeAllKeys(username);
95     }
96
97     protected SshClient getClient() {
98         SshClient client = SshClient.setUpDefaultClient();
99         client.setServerKeyVerifier(new ServerKeyVerifier() {
100             @Override
101             public boolean verifyServerKey(ClientSession sshClientSession, SocketAddress remoteAddress, PublicKey serverKey) {
102                 return true;
103             }
104         });
105         client.start();
106         return client;
107     }
108
109     protected String testSshCommand(String cmd) throws IOException, InterruptedException {
110         return testSshCommand(cmd, null);
111     }
112
113     protected String testSshCommand(String cmd, String stdin) throws IOException, InterruptedException {
114         SshClient client = getClient();
115         ClientSession session = client.connect(username, "localhost", GitBlitSuite.sshPort).await().getSession();
116         session.addPublicKeyIdentity(rwKeyPair);
117         assertTrue(session.auth().await().isSuccess());
118
119         ClientChannel channel = session.createChannel(ClientChannel.CHANNEL_EXEC, cmd);
120         ByteArrayOutputStream baos = new ByteArrayOutputStream();
121         if (stdin != null) {
122             Writer w = new OutputStreamWriter(baos);
123             w.write(stdin);
124             w.close();
125         }
126         channel.setIn(new ByteArrayInputStream(baos.toByteArray()));
127
128         ByteArrayOutputStream out = new ByteArrayOutputStream();
129         ByteArrayOutputStream err = new ByteArrayOutputStream();
130         channel.setOut(out);
131         channel.setErr(err);
132         channel.open();
133
134         channel.waitFor(ClientChannel.CLOSED, 0);
135
136         String result = out.toString().trim();
137         channel.close(false);
138         client.stop();
139         return result;
140     }
141 }