1/*
2 * Copyright (C) 2014 Square, Inc.
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 */
16package com.squareup.okhttp.curl;
17
18import com.google.common.base.Function;
19import com.google.common.base.Joiner;
20import com.google.common.collect.Lists;
21import com.squareup.okhttp.ConnectionPool;
22import com.squareup.okhttp.Headers;
23import com.squareup.okhttp.MediaType;
24import com.squareup.okhttp.OkHttpClient;
25import com.squareup.okhttp.Protocol;
26import com.squareup.okhttp.Request;
27import com.squareup.okhttp.Response;
28import io.airlift.command.Arguments;
29import io.airlift.command.Command;
30import io.airlift.command.HelpOption;
31import io.airlift.command.Option;
32import io.airlift.command.SingleCommand;
33import java.io.IOException;
34import java.io.InputStream;
35import java.security.cert.CertificateException;
36import java.security.cert.X509Certificate;
37import java.util.Arrays;
38import java.util.List;
39import java.util.Properties;
40import javax.net.ssl.SSLContext;
41import javax.net.ssl.SSLSocketFactory;
42import javax.net.ssl.TrustManager;
43import javax.net.ssl.X509TrustManager;
44
45import static java.util.concurrent.TimeUnit.SECONDS;
46
47@Command(name = Main.NAME, description = "A curl for the next-generation web.")
48public class Main extends HelpOption implements Runnable {
49  static final String NAME = "okcurl";
50  static final int DEFAULT_TIMEOUT = -1;
51
52  static Main fromArgs(String... args) {
53    return SingleCommand.singleCommand(Main.class).parse(args);
54  }
55
56  public static void main(String... args) {
57    fromArgs(args).run();
58  }
59
60  private static String versionString() {
61    try {
62      Properties prop = new Properties();
63      InputStream in = Main.class.getResourceAsStream("/okcurl-version.properties");
64      prop.load(in);
65      in.close();
66      return prop.getProperty("version");
67    } catch (IOException e) {
68      throw new AssertionError("Could not load okcurl-version.properties.");
69    }
70  }
71
72  private static String protocols() {
73    return Joiner.on(", ").join(Lists.transform(Arrays.asList(Protocol.values()),
74        new Function<Protocol, String>() {
75          @Override public String apply(Protocol protocol) {
76            return protocol.name.utf8();
77          }
78        }));
79  }
80
81  @Option(name = { "-X", "--request" }, description = "Specify request command to use")
82  public String method;
83
84  @Option(name = { "-d", "--data" }, description = "HTTP POST data")
85  public String data;
86
87  @Option(name = { "-H", "--header" }, description = "Custom header to pass to server")
88  public List<String> headers;
89
90  @Option(name = { "-A", "--user-agent" }, description = "User-Agent to send to server")
91  public String userAgent = NAME + "/" + versionString();
92
93  @Option(name = "--connect-timeout", description = "Maximum time allowed for connection (seconds)")
94  public int connectTimeout = DEFAULT_TIMEOUT;
95
96  @Option(name = "--read-timeout", description = "Maximum time allowed for reading data (seconds)")
97  public int readTimeout = DEFAULT_TIMEOUT;
98
99  @Option(name = { "-L", "--location" }, description = "Follow redirects")
100  public boolean followRedirects;
101
102  @Option(name = { "-k", "--insecure" },
103      description = "Allow connections to SSL sites without certs")
104  public boolean allowInsecure;
105
106  @Option(name = { "-i", "--include" }, description = "Include protocol headers in the output")
107  public boolean showHeaders;
108
109  @Option(name = { "-e", "--referer" }, description = "Referer URL")
110  public String referer;
111
112  @Option(name = { "-V", "--version" }, description = "Show version number and quit")
113  public boolean version;
114
115  @Arguments(title = "url", description = "Remote resource URL")
116  public String url;
117
118  private OkHttpClient client;
119
120  @Override public void run() {
121    if (showHelpIfRequested()) {
122      return;
123    }
124    if (version) {
125      System.out.println(NAME + " " + versionString());
126      System.out.println("Protocols: " + protocols());
127      return;
128    }
129
130    client = createClient();
131    Request request = createRequest();
132    try {
133      Response response = client.execute(request);
134      if (showHeaders) {
135        System.out.println(response.statusLine());
136        Headers headers = response.headers();
137        for (int i = 0, count = headers.size(); i < count; i++) {
138          System.out.println(headers.name(i) + ": " + headers.value(i));
139        }
140        System.out.println();
141      }
142
143      Response.Body body = response.body();
144      byte[] buffer = new byte[1024];
145      while (body.ready()) {
146        int c = body.byteStream().read(buffer);
147        if (c == -1) {
148          return;
149        }
150        System.out.write(buffer, 0, c);
151      }
152      body.close();
153    } catch (IOException e) {
154      e.printStackTrace();
155    } finally {
156      close();
157    }
158  }
159
160  private OkHttpClient createClient() {
161    OkHttpClient client = new OkHttpClient();
162    client.setFollowProtocolRedirects(followRedirects);
163    if (connectTimeout != DEFAULT_TIMEOUT) {
164      client.setConnectTimeout(connectTimeout, SECONDS);
165    }
166    if (readTimeout != DEFAULT_TIMEOUT) {
167      client.setReadTimeout(readTimeout, SECONDS);
168    }
169    if (allowInsecure) {
170      client.setSslSocketFactory(createInsecureSslSocketFactory());
171    }
172    // If we don't set this reference, there's no way to clean shutdown persistent connections.
173    client.setConnectionPool(ConnectionPool.getDefault());
174    return client;
175  }
176
177  private String getRequestMethod() {
178    if (method != null) {
179      return method;
180    }
181    if (data != null) {
182      return "POST";
183    }
184    return "GET";
185  }
186
187  private Request.Body getRequestBody() {
188    if (data == null) {
189      return null;
190    }
191    String bodyData = data;
192
193    String mimeType = "application/x-form-urlencoded";
194    if (headers != null) {
195      for (String header : headers) {
196        String[] parts = header.split(":", -1);
197        if ("Content-Type".equalsIgnoreCase(parts[0])) {
198          mimeType = parts[1].trim();
199          headers.remove(header);
200          break;
201        }
202      }
203    }
204
205    return Request.Body.create(MediaType.parse(mimeType), bodyData);
206  }
207
208  Request createRequest() {
209    Request.Builder request = new Request.Builder();
210
211    request.url(url);
212    request.method(getRequestMethod(), getRequestBody());
213
214    if (headers != null) {
215      for (String header : headers) {
216        String[] parts = header.split(":", -1);
217        request.header(parts[0], parts[1]);
218      }
219    }
220    if (referer != null) {
221      request.header("Referer", referer);
222    }
223    request.header("User-Agent", userAgent);
224
225    return request.build();
226  }
227
228  private void close() {
229    client.getConnectionPool().evictAll(); // Close any persistent connections.
230  }
231
232  private static SSLSocketFactory createInsecureSslSocketFactory() {
233    try {
234      SSLContext context = SSLContext.getInstance("TLS");
235      TrustManager permissive = new X509TrustManager() {
236        @Override public void checkClientTrusted(X509Certificate[] chain, String authType)
237            throws CertificateException {
238        }
239
240        @Override public void checkServerTrusted(X509Certificate[] chain, String authType)
241            throws CertificateException {
242        }
243
244        @Override public X509Certificate[] getAcceptedIssuers() {
245          return null;
246        }
247      };
248      context.init(null, new TrustManager[] { permissive }, null);
249      return context.getSocketFactory();
250    } catch (Exception e) {
251      throw new AssertionError(e);
252    }
253  }
254}
255