From 19363ee46d9dd72d2746aec60970c02bc6923ed0 Mon Sep 17 00:00:00 2001
From: Jimmy Hedman <jimmy.hedman@gmail.com>
Date: Mon, 21 May 2018 14:30:36 +0200
Subject: [PATCH] Add domain paramteter to wifi.

- To be able to do OTA updates on networks that doesn't use .local as
  local domain parameter domain is added to wifi section. It's currently
  only used for OTA.
---
 esphomeyaml/__main__.py        | 9 +++++----
 esphomeyaml/components/wifi.py | 7 ++++---
 esphomeyaml/const.py           | 1 +
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/esphomeyaml/__main__.py b/esphomeyaml/__main__.py
index c79da1971b..9934522b6a 100644
--- a/esphomeyaml/__main__.py
+++ b/esphomeyaml/__main__.py
@@ -8,8 +8,8 @@ import sys
 
 from esphomeyaml import core, mqtt, wizard, writer, yaml_util, const
 from esphomeyaml.config import core_to_code, get_component, iter_components, read_config
-from esphomeyaml.const import CONF_BAUD_RATE, CONF_ESPHOMEYAML, CONF_HOSTNAME, CONF_LOGGER, \
-    CONF_MANUAL_IP, CONF_NAME, CONF_STATIC_IP, CONF_WIFI
+from esphomeyaml.const import CONF_BAUD_RATE, CONF_DOMAIN, CONF_ESPHOMEYAML, CONF_HOSTNAME, \
+    CONF_LOGGER, CONF_MANUAL_IP, CONF_NAME, CONF_STATIC_IP, CONF_WIFI
 from esphomeyaml.core import ESPHomeYAMLError
 from esphomeyaml.helpers import AssignmentExpression, RawStatement, _EXPRESSIONS, add, add_task, \
     color, get_variable, indent, quote, statement, Expression
@@ -163,12 +163,13 @@ def upload_program(config, args, port):
         _LOGGER.error("No serial port found and OTA not enabled. Can't upload!")
         return -1
 
+    DOMAIN = '.' + config[CONF_WIFI].get(CONF_DOMAIN, 'local')
     if CONF_MANUAL_IP in config[CONF_WIFI]:
         host = str(config[CONF_WIFI][CONF_MANUAL_IP][CONF_STATIC_IP])
     elif CONF_HOSTNAME in config[CONF_WIFI]:
-        host = config[CONF_WIFI][CONF_HOSTNAME] + u'.local'
+        host = config[CONF_WIFI][CONF_HOSTNAME] + DOMAIN
     else:
-        host = config[CONF_ESPHOMEYAML][CONF_NAME] + u'.local'
+        host = config[CONF_ESPHOMEYAML][CONF_NAME] + DOMAIN
 
     from esphomeyaml.components import ota
     from esphomeyaml import espota
diff --git a/esphomeyaml/components/wifi.py b/esphomeyaml/components/wifi.py
index c1d09010fe..5b74430846 100644
--- a/esphomeyaml/components/wifi.py
+++ b/esphomeyaml/components/wifi.py
@@ -2,9 +2,9 @@ import voluptuous as vol
 
 import esphomeyaml.config_validation as cv
 from esphomeyaml import core
-from esphomeyaml.const import CONF_AP, CONF_CHANNEL, CONF_DNS1, CONF_DNS2, CONF_GATEWAY, \
-    CONF_HOSTNAME, CONF_ID, CONF_MANUAL_IP, CONF_PASSWORD, CONF_SSID, CONF_STATIC_IP, CONF_SUBNET, \
-    ESP_PLATFORM_ESP8266
+from esphomeyaml.const import CONF_AP, CONF_CHANNEL, CONF_DNS1, CONF_DNS2, CONF_DOMAIN, \
+    CONF_GATEWAY, CONF_HOSTNAME, CONF_ID, CONF_MANUAL_IP, CONF_PASSWORD, CONF_SSID, \
+    CONF_STATIC_IP, CONF_SUBNET, ESP_PLATFORM_ESP8266
 from esphomeyaml.helpers import App, Pvariable, StructInitializer, add, esphomelib_ns, global_ns
 
 
@@ -42,6 +42,7 @@ CONFIG_SCHEMA = vol.Schema({
         vol.Optional(CONF_MANUAL_IP): AP_MANUAL_IP_SCHEMA,
     }),
     vol.Optional(CONF_HOSTNAME): cv.hostname,
+    vol.Optional(CONF_DOMAIN): cv.string,
 })
 
 # pylint: disable=invalid-name
diff --git a/esphomeyaml/const.py b/esphomeyaml/const.py
index 685019b088..7e7395f1e8 100644
--- a/esphomeyaml/const.py
+++ b/esphomeyaml/const.py
@@ -228,6 +228,7 @@ CONF_TURN_OFF_ACTION = 'turn_off_action'
 CONF_OPEN_ACTION = 'open_action'
 CONF_CLOSE_ACTION = 'close_action'
 CONF_STOP_ACTION = 'stop_action'
+CONF_DOMAIN = 'domain'
 
 ESP32_BOARDS = [
     'featheresp32', 'node32s', 'espea32', 'firebeetle32', 'esp32doit-devkit-v1',